From ce7ebbc9fba562c8d56ca1b4a5aa15772fd0eb4f Mon Sep 17 00:00:00 2001
From: kach <hardmath123@gmail.com>
Date: Sat, 17 Aug 2024 17:48:53 -0400
Subject: [PATCH] do not require JAX for build...

---
 memo/codegen.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/memo/codegen.py b/memo/codegen.py
index 75a0d64..a96c3a6 100644
--- a/memo/codegen.py
+++ b/memo/codegen.py
@@ -5,7 +5,6 @@
 import textwrap
 import os, sys, platform, inspect
 from io import StringIO
-import jax
 from typing import Any
 
 def codegen(
@@ -127,6 +126,9 @@ def memo(f=None, **kwargs):  # type: ignore
                 "[We think this may be a bug in memo: if you don't understand what is going on, please get in touch with us!]"
             )
         e.add_note("")
+
+        # Describe environment...
+        import jax
         e.add_note(f"  P.S.: You are currently using...")
         e.add_note(f"        + memo version {__version__} and JAX version {jax.__version__}")
         e.add_note(f"        + on Python version {platform.python_version()} on the {platform.system()} platform")