Fixes #188.

patrick-kidger · Mar 7, 2024 · 102e499 · 102e499
1 parent f708d17
commit 102e499
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
@@ -189,7 +189,11 @@ def f(...): ...
     """
 
     global _tb_flag
-    if _tb_flag and importlib.util.find_spec("jax._src.traceback_util") is not None:
+    if (
+        _tb_flag
+        and importlib.util.find_spec("jax") is not None
+        and importlib.util.find_spec("jax._src.traceback_util") is not None
+    ):
         import jax._src.traceback_util as traceback_util
 
         traceback_util.register_exclusion(__file__)

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.27"
+version = "0.2.28"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"