Fixes #188.

jaxtyping/__init__.py

@@ -145,6 +145,7 @@ else:
                         pass
 
                     Array.__module__ = "builtins"
+                    Array.__qualname__ = "Array"
                     return Array
                 else:
                     import jax
@@ -157,6 +158,7 @@ else:
                         pass
 
                     ArrayLike.__module__ = "builtins"
+                    ArrayLike.__qualname__ = "ArrayLike"
                     return ArrayLike
                 else:
                     import jax.typing
@@ -169,6 +171,7 @@ else:
                         pass
 
                     PRNGKeyArray.__module__ = "builtins"
+                    PRNGKeyArray.__qualname__ = "PRNGKeyArray"
                     return PRNGKeyArray
                 else:
                     # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
@@ -216,6 +219,7 @@ else:
                     if typing.GENERATING_DOCUMENTATION:
                         # Equinox etc. docs get just `PyTreeDef`.
                         # jaxtyping docs get `jaxtyping.PyTreeDef`.
+                        PyTreeDef.__qualname__ = "PyTreeDef"
                         PyTreeDef.__module__ = "builtins"
                     return PyTreeDef
                 else:

jaxtyping/_array_types.py

@@ -229,7 +229,7 @@ class _MetaAbstractArray(type):
             return _check_dims(cls.dims, obj.shape, single_memo, arg_memo)
         else:
             if obj.ndim < len(cls.dims) - 1:
-                return f"this array has {obj.ndim} dimensions, which is fewer than {len(cls.dims - 1)} that is the minimum expected by the type hint"  # noqa: E501
+                return f"this array has {obj.ndim} dimensions, which is fewer than {len(cls.dims) - 1} that is the minimum expected by the type hint"  # noqa: E501
             i = cls.index_variadic
             j = -(len(cls.dims) - i - 1)
             if j == 0:

jaxtyping/_decorator.py

@@ -189,7 +189,11 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
     """
 
     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__)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.26"
+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"