link to why use type annotations

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

README.md

@@ -1,6 +1,6 @@
 <h1 align="center">jaxtyping</h1>
 
-Type annotations **and runtime type-checking** for:
+[Use type annotations **and runtime type-checking**](https://jax.readthedocs.io/en/latest/jep/12049-type-annotations.html) for:
 
 1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
@@ -8,7 +8,11 @@ Type annotations **and runtime type-checking** for:
 
 **For example:**
 ```python
-from jaxtyping import Array, Float, PyTree
+from jaxtyping import Array, Float, PyTree, , UInt, Int, Bool
+import torch
+impport numpy as np
+import tensorflow as tf
+
 
 # Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
@@ -21,6 +25,15 @@ def accepts_pytree_of_ints(x: PyTree[int]):
 
 def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
     ...
+
+def accepts_torch.Long(x: Int[torch.Tensor, "batch channel height width"]):
+  ....
+
+def accepts_numpy_float(x :Float[np.ndarray, "batch sequence features"]):
+  ...
+
+def accepts_tensorflow_uint(x: hint = UInt[tf.Tensor, "b c h w"]):
+  ...
 ```
 
 ## Installation
@@ -39,32 +52,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.  
-
-[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).  
 
+**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!)  
 
-### Disclaimer
+**Awesome JAX**  
-
+[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  
-This is not an official Google product.

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!  
 
+**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).  
-[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
-
-[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).  
+
+**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!)  
+
+**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):
+            if len(obj.shape) != len(cls.dims):
-                return f"this array has {obj.ndim} dimensions, not the {len(cls.dims)} expected by the type hint"  # noqa: E501
+                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:
+            if len(obj.shape) < 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 {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:
+    return ChooseTypecheckerMagics
-    # Very broad exception-handling, as e.g. IPython will sometimes be
-    # present but fail to import for mysterious reasons.
-    pass
 
 
 def load_ipython_extension(ipython):
     try:
-        ipython.register_magics(ChooseTypecheckerMagics)
+        ChooseTypecheckerMagics = choose_typechecker_magics()
-    except NameError:
+    except Exception as e:
-        raise NameError(
+        # Very broad exception-handling, as e.g. IPython will sometimes be
-            "ChooseTypecheckerMagics is not defined.\n\n"
+        # present but fail to import for mysterious reasons.
-            + "You may be trying to use IPython extension without IPython installed."
+        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"