also require jaxlib

patrick-kidger · Mar 8, 2024 · 696cc5b · 696cc5b
1 parent 102e499
commit 696cc5b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jaxtyping/_decorator.py b/jaxtyping/_decorator.py
@@ -192,6 +192,7 @@ def f(...): ...
     if (
         _tb_flag
         and importlib.util.find_spec("jax") is not None
+        and importlib.util.find_spec("jaxlib") is not None
         and importlib.util.find_spec("jax._src.traceback_util") is not None
     ):
         import jax._src.traceback_util as traceback_util