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")