Updated ecosystem again

IPython is quite a chunky package, and importing it unconditionally makes
jaxtyping itself relatively slow to import.

README.md

@@ -39,32 +39,23 @@ The annotations provided by jaxtyping are compatible with runtime type-checking
 
 Available at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxtyping).
 
-## Finally
+## See also: other libraries in the JAX ecosystem
 
-### See also: other libraries in the JAX ecosystem
+**Always useful**  
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks and everything not already in core JAX!  
 
-[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+**Deep learning**  
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  
+[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).  
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).  
 
-[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
+**Scientific computing**  
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  
+[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  
+[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.  
+[sympy2jax](https://github.com/patrick-kidger/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.  
+[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  
 
-[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
-
-[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.
-
-[Lineax](https://github.com/google/lineax): linear solvers.
-
-[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.
-
-[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).
-
-[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
-
-[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
-
-[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
-
-[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)
-
-### Disclaimer
-
-This is not an official Google product.
+**Awesome JAX**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  

docs/index.md

@@ -43,16 +43,21 @@ Have a read of the [Array annotations](./api/array.md) documentation on the left
 
 ## See also: other libraries in the JAX ecosystem
 
-[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+**Always useful**  
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks and everything not already in core JAX!  
 
-[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
+**Deep learning**  
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.  
+[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).  
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).  
 
-[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
+**Scientific computing**  
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.  
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.  
+[Lineax](https://github.com/patrick-kidger/lineax): linear solvers.  
+[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.  
+[sympy2jax](https://github.com/patrick-kidger/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.  
+[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  
 
-[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
-
-[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
-
-[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
-
-[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
+**Awesome JAX**  
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  

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

@@ -224,12 +224,12 @@ class _MetaAbstractArray(type):
         arg_memo: dict[str, Any],
     ) -> str:
         if cls.index_variadic is None:
-            if obj.ndim != len(cls.dims):
-                return f"this array has {obj.ndim} dimensions, not the {len(cls.dims)} expected by the type hint"  # noqa: E501
+            if len(obj.shape) != len(cls.dims):
+                return f"this array has {len(obj.shape)} dimensions, not the {len(cls.dims)} expected by the type hint"  # noqa: E501
             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
+            if len(obj.shape) < len(cls.dims) - 1:
+                return f"this array has {len(obj.shape)} 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,12 @@ 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("jaxlib") 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__)

jaxtyping/_ipython_extension.py

@@ -20,7 +20,9 @@
 from ._import_hook import JaxtypingTransformer, Typechecker
 
 
-try:
+def choose_typechecker_magics():
+    # The import is local to avoid degrading import times when the magic is
+    # not needed.
     from IPython.core.magic import line_magic, Magics, magics_class
 
     @magics_class
@@ -40,17 +42,15 @@ try:
                 JaxtypingTransformer(typechecker=Typechecker(typechecker))
             )
 
-except Exception:
-    # Very broad exception-handling, as e.g. IPython will sometimes be
-    # present but fail to import for mysterious reasons.
-    pass
+    return ChooseTypecheckerMagics
 
 
 def load_ipython_extension(ipython):
     try:
-        ipython.register_magics(ChooseTypecheckerMagics)
-    except NameError:
-        raise NameError(
-            "ChooseTypecheckerMagics is not defined.\n\n"
-            + "You may be trying to use IPython extension without IPython installed."
-        )
+        ChooseTypecheckerMagics = choose_typechecker_magics()
+    except Exception as e:
+        # Very broad exception-handling, as e.g. IPython will sometimes be
+        # present but fail to import for mysterious reasons.
+        raise RuntimeError("Failed to define jaxtyping.typechecker magic") from e
+
+    ipython.register_magics(ChooseTypecheckerMagics)

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"