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/_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

@@ -192,6 +192,7 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
     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

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)