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/api/advanced-features.md

@@ -7,6 +7,10 @@
         members:
             false
 
+## Printing axis bindings
+
+::: jaxtyping.print_bindings
+
 ## Introspection
 
 If you're writing your own type hint parser, then you may wish to detect if some Python object is a jaxtyping-provided type.

docs/api/array.md

@@ -25,7 +25,7 @@ In addition some modifiers can be applied:
     `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
 - Prepend `_` to an axis to disable any runtime checking of that axis (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
 - Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
-- Prepend `?` to an axis to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree/).)
+- Prepend `?` to an axis to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](./pytree.md).)
 
 When using multiple modifiers, their order does not matter.
 
@@ -66,9 +66,9 @@ The dtype should be any one of (all imported from `jaxtyping`):
                 - Of particular precision: `Complex64`, `Complex128`
         - Any integer or unsigned intger: `Integer`
             - Any unsigned integer: `UInt`
-                - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
+                - Of particular precision: `UInt4`, `UInt8`, `UInt16`, `UInt32`, `UInt64`
             - Any signed integer: `Int`
-                - Of particular precision: `Int8`, `Int16`, `Int32`, `Int64`
+                - Of particular precision: `Int4`, `Int8`, `Int16`, `Int32`, `Int64`
         - Any floating, integer, or unsigned integer: `Real`.
 
 Unless you really want to force a particular precision, then for most applications you should probably allow any floating-point, any integer, etc. That is, use

docs/api/runtime-type-checking.md

@@ -8,6 +8,10 @@ There are two approaches: either use [`jaxtyping.jaxtyped`][] to typecheck a sin
 
 In either case, the actual business of checking types is performed with the help of a runtime type-checking library. The two most popular are [beartype](https://github.com/beartype/beartype) and [typeguard](https://github.com/agronholm/typeguard). (If using typeguard, then specifically the version `2.*` series should be used. Later versions -- `3` and `4` -- have some known issues.)
 
+!!! warning
+
+    Avoid using `from __future__ import annotations`, or stringified type annotations, where possible. These are largely incompatible with runtime type checking. See also [this FAQ entry](../faq.md#dataclass-annotations-arent-being-checked-properly).
+
 ---
 
 ::: jaxtyping.jaxtyped

docs/faq.md

@@ -20,6 +20,24 @@ Some tooling in the Python ecosystem assumes that only the latter is true, and w
 
 In the case of `flake8`, or Ruff, this can be resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
 
+## Dataclass annotations aren't being checked properly.
+
+Stringified dataclass annotations, e.g.
+```python
+@dataclass()
+class Foo:
+   x: "int"
+```
+will be silently skipped without checking them. This is because these are essentially impossible to resolve at runtime. Such stringified annotations typically occur either when using them for forward references, or when using `from __future__ import annotations`. (You should essentially never use the latter, it is largely incompatible with runtime type checking and as such is [being replaced in Python 3.13](https://peps.python.org/pep-0649/).)
+
+Partially stringified dataclass annotations, e.g.
+```python
+@dataclass()
+class Foo:
+   x: tuple["int"]
+```
+will likely raise an error, and must not be used at all.
+
 ## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)?
 
 The intention of PEP 646 was to make it possible for static type checkers to perform shape checks of arrays. Unfortunately, this still isn't yet practical, so jaxtyping deliberately does not use this. (Yet?)

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

@@ -17,16 +17,17 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+import functools as ft
 import importlib.metadata
+import importlib.util
 import typing
 import warnings
+from typing import Union
 
-# First import some things as normal
 from ._array_types import (
     AbstractArray as AbstractArray,
     AbstractDtype as AbstractDtype,
     get_array_name_format as get_array_name_format,
-    has_jax,
     set_array_name_format as set_array_name_format,
 )
 from ._config import config as config
@@ -37,41 +38,18 @@ from ._errors import (
 )
 from ._import_hook import install_import_hook as install_import_hook
 from ._ipython_extension import load_ipython_extension as load_ipython_extension
+from ._storage import print_bindings as print_bindings
 
 
-# Now import Array and ArrayLike
 if typing.TYPE_CHECKING:
-    # For imports, we need to explicitly `import X as X` in order for Pyright to see
-    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+    import typing_extensions
+
     from jax import Array as Array
-    from jax.typing import ArrayLike as ArrayLike
-elif has_jax:
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    from jax.tree_util import PyTreeDef as PyTreeDef
+    from jax.typing import ArrayLike as ArrayLike, DTypeLike as DTypeLike
 
-        class Array:
-            pass
-
-        Array.__module__ = "builtins"
-
-        class ArrayLike:
-            pass
-
-        ArrayLike.__module__ = "builtins"
-    else:
-        from jax import Array as Array
-
-        try:
-            from jax.typing import ArrayLike as ArrayLike
-        except (ModuleNotFoundError, ImportError):
-            pass
-
-
-# Import our dtypes
-if typing.TYPE_CHECKING:
     # Introduce an indirection so that we can `import X as X` to make it clear that
     # these are public.
-    from jax.typing import DTypeLike as DTypeLike
-
     from ._indirection import (
         BFloat16 as BFloat16,
         Bool as Bool,
@@ -84,6 +62,7 @@ if typing.TYPE_CHECKING:
         Float64 as Float64,
         Inexact as Inexact,
         Int as Int,
+        Int4 as Int4,
         Int8 as Int8,
         Int16 as Int16,
         Int32 as Int32,
@@ -91,56 +70,18 @@ if typing.TYPE_CHECKING:
         Integer as Integer,
         Key as Key,
         Num as Num,
+        PRNGKeyArray as PRNGKeyArray,
         Real as Real,
+        Scalar as Scalar,
+        ScalarLike as ScalarLike,
         Shaped as Shaped,
         UInt as UInt,
+        UInt4 as UInt4,
         UInt8 as UInt8,
         UInt16 as UInt16,
         UInt32 as UInt32,
         UInt64 as UInt64,
     )
-else:
-    from ._array_types import (
-        BFloat16 as BFloat16,
-        Bool as Bool,
-        Complex as Complex,
-        Complex64 as Complex64,
-        Complex128 as Complex128,
-        Float as Float,
-        Float16 as Float16,
-        Float32 as Float32,
-        Float64 as Float64,
-        Inexact as Inexact,
-        Int as Int,
-        Int8 as Int8,
-        Int16 as Int16,
-        Int32 as Int32,
-        Int64 as Int64,
-        Integer as Integer,
-        Num as Num,
-        Real as Real,
-        Shaped as Shaped,
-        UInt as UInt,
-        UInt8 as UInt8,
-        UInt16 as UInt16,
-        UInt32 as UInt32,
-        UInt64 as UInt64,
-    )
-
-    if has_jax:
-        import jax.typing
-
-        from ._array_types import Key as Key
-
-        if hasattr(jax.typing, "DTypeLike"):
-            from jax.typing import DTypeLike as DTypeLike
-
-
-# Now import PyTreeDef and PyTree
-if typing.TYPE_CHECKING:
-    import typing_extensions
-
-    from jax.tree_util import PyTreeDef as PyTreeDef
 
     # Set up to deliberately confuse a static type checker.
     PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar")
@@ -160,57 +101,133 @@ if typing.TYPE_CHECKING:
     # If they can't figure out what a type is, then they just give up and allow
     # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
     # annotation, which static type checkers aren't smart enough to resolve.
-elif has_jax:
-    if hasattr(typing, "GENERATING_DOCUMENTATION"):
-        # Most parts of the Equinox ecosystem have
-        # `typing.GENERATING_DOCUMENTATION = True` when generating documentation, to
-        # add whatever shims are necessary to get pretty docs. E.g. to have type
-        # annotations appear as just `PyTree`, not `jaxtyping.PyTree`.
-        #
-        # As jaxtyping actually wants things to appear as e.g. `jaxtyping.PyTree`,
-        # rather than just `PyTree`, then it sets
-        # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims.
-        #
-        # Here we do only a `hasattr` check, as we want to get this version of
-        # `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs.
-
-        class PyTreeDef:
-            """Alias for `jax.tree_util.PyTreeDef`, which is the type of the return
-            from `jax.tree_util.tree_structure(...)`.
-            """
-
-        if typing.GENERATING_DOCUMENTATION:
-            # Equinox etc. docs get just `PyTreeDef`.
-            # jaxtyping docs get `jaxtyping.PyTreeDef`.
-            PyTreeDef.__module__ = "builtins"
-
-    else:
-        from jax.tree_util import PyTreeDef as PyTreeDef
-
-    from ._pytree_type import PyTree as PyTree  # noqa: F401
-
-
-# Conveniences
-if typing.TYPE_CHECKING:
-    from ._indirection import (
-        PRNGKeyArray as PRNGKeyArray,
-        Scalar as Scalar,
-        ScalarLike as ScalarLike,
+else:
+    from ._array_types import (
+        BFloat16 as BFloat16,
+        Bool as Bool,
+        Complex as Complex,
+        Complex64 as Complex64,
+        Complex128 as Complex128,
+        Float as Float,
+        Float16 as Float16,
+        Float32 as Float32,
+        Float64 as Float64,
+        Inexact as Inexact,
+        Int as Int,
+        Int4 as Int4,
+        Int8 as Int8,
+        Int16 as Int16,
+        Int32 as Int32,
+        Int64 as Int64,
+        Integer as Integer,
+        Key as Key,
+        Num as Num,
+        Real as Real,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt4 as UInt4,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
     )
-elif has_jax:
-    from ._array_types import Scalar, ScalarLike  # noqa: F401
 
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
-        # That is, we're generating some downstream documentation, not the jaxtyping
-        # documentation itself.
-        class PRNGKeyArray:
-            pass
+    # But crucially, does not actually import jax at all. We do that dynamically in
+    # __getattr__ if required. See #178.
+    if importlib.util.find_spec("jax") is not None:
 
-        PRNGKeyArray.__module__ = "builtins"
-    else:
-        from ._array_types import PRNGKeyArray
+        @ft.cache
+        def __getattr__(item):
+            if item == "Array":
+                if getattr(typing, "GENERATING_DOCUMENTATION", False):
 
-del has_jax
+                    class Array:
+                        pass
+
+                    Array.__module__ = "builtins"
+                    Array.__qualname__ = "Array"
+                    return Array
+                else:
+                    import jax
+
+                    return jax.Array
+            elif item == "ArrayLike":
+                if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+                    class ArrayLike:
+                        pass
+
+                    ArrayLike.__module__ = "builtins"
+                    ArrayLike.__qualname__ = "ArrayLike"
+                    return ArrayLike
+                else:
+                    import jax.typing
+
+                    return jax.typing.ArrayLike
+            elif item == "PRNGKeyArray":
+                if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+                    class PRNGKeyArray:
+                        pass
+
+                    PRNGKeyArray.__module__ = "builtins"
+                    PRNGKeyArray.__qualname__ = "PRNGKeyArray"
+                    return PRNGKeyArray
+                else:
+                    # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
+                    # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype
+                    # `uint32`.
+                    import jax
+
+                    return Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
+            elif item == "DTypeLike":
+                import jax.typing
+
+                return jax.typing.DTypeLike
+            elif item == "Scalar":
+                import jax
+
+                return Shaped[jax.Array, ""]
+            elif item == "ScalarLike":
+                import jax.typing
+
+                return Shaped[jax.typing.ArrayLike, ""]
+            elif item == "PyTree":
+                from ._pytree_type import PyTree
+
+                return PyTree
+            elif item == "PyTreeDef":
+                if hasattr(typing, "GENERATING_DOCUMENTATION"):
+                    # Most parts of the Equinox ecosystem have
+                    # `typing.GENERATING_DOCUMENTATION = True` when generating
+                    # documentation, to add whatever shims are necessary to get pretty
+                    # docs. E.g. to have type annotations appear as just `PyTree`, not
+                    # `jaxtyping.PyTree`.
+                    #
+                    # As jaxtyping actually wants things to appear as e.g.
+                    # `jaxtyping.PyTree`, rather than just `PyTree`, then it sets
+                    # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims.
+                    #
+                    # Here we do only a `hasattr` check, as we want to get this version
+                    # of `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs.
+
+                    class PyTreeDef:
+                        """Alias for `jax.tree_util.PyTreeDef`, which is the type of the
+                        return from `jax.tree_util.tree_structure(...)`.
+                        """
+
+                    if typing.GENERATING_DOCUMENTATION:
+                        # Equinox etc. docs get just `PyTreeDef`.
+                        # jaxtyping docs get `jaxtyping.PyTreeDef`.
+                        PyTreeDef.__qualname__ = "PyTreeDef"
+                        PyTreeDef.__module__ = "builtins"
+                    return PyTreeDef
+                else:
+                    import jax.tree_util
+
+                    return jax.tree_util.PyTreeDef
+            else:
+                raise AttributeError(f"module jaxtyping has no attribute {item!r}")
 
 
 check_equinox_version = True  # easy-to-replace line with copybara

jaxtyping/_array_types.py

@@ -23,6 +23,7 @@ import re
 import sys
 import types
 import typing
+from dataclasses import dataclass
 from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
@@ -36,18 +37,6 @@ from ._storage import (
 )
 
 
-try:
-    import jax
-except (ImportError, RuntimeError, AttributeError):
-    # We catch `RuntimeError` as JAX will throw this if it's present, but unable to run
-    # on the current machine. This fails with this error.
-    # We catch `AttributeError` as the above then leaves the module in a partially
-    # initialised state, which causes subsequent imports to fail with this error.
-    has_jax = False
-else:
-    has_jax = True
-
-
 _array_name_format = "dtype_and_shape"
 
 
@@ -62,7 +51,6 @@ def set_array_name_format(value):
 
 _any_dtype = object()
 
-
 _anonymous_dim = object()
 _anonymous_variadic_dim = object()
 
@@ -73,30 +61,30 @@ class _DimType(enum.Enum):
     symbolic = enum.auto()
 
 
+@dataclass(frozen=True)
 class _NamedDim:
-    def __init__(self, name, broadcastable, treepath):
-        self.name = name
-        self.broadcastable = broadcastable
-        self.treepath = treepath
+    name: str
+    broadcastable: bool
+    treepath: Any
 
 
+@dataclass(frozen=True)
 class _NamedVariadicDim:
-    def __init__(self, name, broadcastable, treepath):
-        self.name = name
-        self.broadcastable = broadcastable
-        self.treepath = treepath
+    name: str
+    broadcastable: bool
+    treepath: Any
 
 
+@dataclass(frozen=True)
 class _FixedDim:
-    def __init__(self, size, broadcastable):
-        self.size = size
-        self.broadcastable = broadcastable
+    size: str
+    broadcastable: bool
 
 
+@dataclass(frozen=True)
 class _SymbolicDim:
-    def __init__(self, elem, broadcastable):
-        self.elem = elem
-        self.broadcastable = broadcastable
+    elem: Any
+    broadcastable: bool
 
 
 _AbstractDimOrVariadicDim = Union[
@@ -115,7 +103,7 @@ def _check_dims(
     obj_shape: tuple[int, ...],
     single_memo: dict[str, int],
     arg_memo: dict[str, Any],
-) -> bool:
+) -> str:
     assert len(cls_dims) == len(obj_shape)
     for cls_dim, obj_size in zip(cls_dims, obj_shape):
         if cls_dim is _anonymous_dim:
@@ -124,7 +112,7 @@ def _check_dims(
             pass
         elif type(cls_dim) is _FixedDim:
             if cls_dim.size != obj_size:
-                return False
+                return f"the dimension size {obj_size} does not equal {cls_dim.size} as expected by the type hint"  # noqa: E501
         elif type(cls_dim) is _SymbolicDim:
             try:
                 # Support f-string syntax.
@@ -141,7 +129,7 @@ def _check_dims(
                     "arguments."
                 ) from e
             if eval_size != obj_size:
-                return False
+                return f"the dimension size {obj_size} does not equal the existing value of {cls_dim.elem}={eval_size}"  # noqa: E501
         else:
             assert type(cls_dim) is _NamedDim
             if cls_dim.treepath:
@@ -154,16 +142,26 @@ def _check_dims(
                 single_memo[name] = obj_size
             else:
                 if cls_size != obj_size:
-                    return False
-    return True
+                    return f"the size of dimension {cls_dim.name} is {obj_size} which does not equal the existing value of {cls_size}"  # noqa: E501
+    return ""
 
 
 class _MetaAbstractArray(type):
-    def __instancecheck__(cls, obj):
+    _skip_instancecheck: bool = False
+
+    def make_transparent(cls):
+        cls._skip_instancecheck = True
+
+    def __instancecheck__(cls, obj: Any) -> bool:
+        return cls.__instancecheck_str__(obj) == ""
+
+    def __instancecheck_str__(cls, obj: Any) -> str:
+        if cls._skip_instancecheck:
+            return ""
         if not isinstance(obj, cls.array_type):
-            return False
+            return f"this value is not an instance of the underlying array type {cls.array_type}"  # noqa: E501
         if get_treeflatten_memo():
-            return True
+            return ""
 
         if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
@@ -193,7 +191,10 @@ class _MetaAbstractArray(type):
                 if in_dtypes:
                     break
             if not in_dtypes:
-                return False
+                if len(cls.dtypes) == 1:
+                    return f"this array has dtype {dtype}, not {cls.dtypes[0]} as expected by the type hint"  # noqa: E501
+                else:
+                    return f"this array has dtype {dtype}, not any of {cls.dtypes} as expected by the type hint"  # noqa: E501
 
         single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo()
         single_memo_bak = single_memo.copy()
@@ -207,13 +208,13 @@ class _MetaAbstractArray(type):
                 single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
             )
             raise
-        if check:
-            return True
+        if check == "":
+            return check
         else:
             set_shape_memo(
                 single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
             )
-            return False
+            return check
 
     def _check_shape(
         cls,
@@ -221,27 +222,32 @@ class _MetaAbstractArray(type):
         single_memo: dict[str, int],
         variadic_memo: dict[str, tuple[bool, tuple[int, ...]]],
         arg_memo: dict[str, Any],
-    ):
+    ) -> str:
         if cls.index_variadic is None:
-            if obj.ndim != len(cls.dims):
-                return False
+            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 False
+            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:
                 j = None
-            if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo, arg_memo):
-                return False
-            if j is not None and not _check_dims(
-                cls.dims[j:], obj.shape[j:], single_memo, arg_memo
-            ):
-                return False
+            prefix_check = _check_dims(
+                cls.dims[:i], obj.shape[:i], single_memo, arg_memo
+            )
+            if prefix_check != "":
+                return prefix_check
+            if j is not None:
+                suffix_check = _check_dims(
+                    cls.dims[j:], obj.shape[j:], single_memo, arg_memo
+                )
+                if suffix_check != "":
+                    return suffix_check
             variadic_dim = cls.dims[i]
             if variadic_dim is _anonymous_variadic_dim:
-                return True
+                return ""
             else:
                 assert type(variadic_dim) is _NamedVariadicDim
                 if variadic_dim.treepath:
@@ -253,16 +259,16 @@ class _MetaAbstractArray(type):
                     prev_broadcastable, prev_shape = variadic_memo[name]
                 except KeyError:
                     variadic_memo[name] = (broadcastable, obj.shape[i:j])
-                    return True
+                    return ""
                 else:
                     new_shape = obj.shape[i:j]
                     if prev_broadcastable:
                         try:
                             broadcast_shape = np.broadcast_shapes(new_shape, prev_shape)
                         except ValueError:  # not broadcastable e.g. (3, 4) and (5,)
-                            return False
+                            return f"the shape of its variadic dimensions '*{variadic_dim.name}' is {new_shape}, which cannot be broadcast with the existing value of {prev_shape}"  # noqa: E501
                         if not broadcastable and broadcast_shape != new_shape:
-                            return False
+                            return f"the shape of its variadic dimensions '*{variadic_dim.name}' is {new_shape}, which the existing value of {prev_shape} cannot be broadcast to"  # noqa: E501
                         variadic_memo[name] = (broadcastable, broadcast_shape)
                     else:
                         if broadcastable:
@@ -271,20 +277,37 @@ class _MetaAbstractArray(type):
                                     new_shape, prev_shape
                                 )
                             except ValueError:  # not broadcastable e.g. (3, 4) and (5,)
-                                return False
+                                return f"the shape of its variadic dimensions '*{variadic_dim.name}' is {new_shape}, which cannot be broadcast with the existing value of {prev_shape}"  # noqa: E501
                             if broadcast_shape != prev_shape:
-                                return False
+                                return f"the shape of its variadic dimensions '*{variadic_dim.name}' is {new_shape}, which cannot be broadcast to the existing value of {prev_shape}"  # noqa: E501
                         else:
                             if new_shape != prev_shape:
-                                return False
-                    return True
+                                return f"the shape of its variadic dimensions '*{variadic_dim.name}' is {new_shape}, which does not equal the existing value of {prev_shape}"  # noqa: E501
+                    return ""
             assert False
 
 
 @ft.lru_cache(maxsize=None)
 def _make_metaclass(base_metaclass):
     class MetaAbstractArray(_MetaAbstractArray, base_metaclass):
-        pass
+        def _get_props(cls):
+            props_tuple = (
+                cls.index_variadic,
+                cls.dims,
+                cls.array_type,
+                cls.dtypes,
+                cls.dim_str,
+            )
+            return props_tuple
+
+        def __eq__(cls, other):
+            if type(cls) is not type(other):
+                return False
+
+            return cls._get_props() == other._get_props()
+
+        def __hash__(cls):
+            return hash(cls._get_props())
 
     return MetaAbstractArray
 
@@ -315,14 +338,13 @@ class AbstractArray(metaclass=_MetaAbstractArray):
 
 _not_made = object()
 
-
 _union_types = [typing.Union]
 if sys.version_info >= (3, 10):
     _union_types.append(types.UnionType)
 
 
 @ft.lru_cache(maxsize=None)
-def _make_array(array_type, dim_str, dtypes, name):
+def _make_array_cached(array_type, dim_str, dtypes, name):
     if not isinstance(dim_str, str):
         raise ValueError(
             "Shape specification must be a string. Axes should be separated with "
@@ -537,22 +559,33 @@ def _make_array(array_type, dim_str, dtypes, name):
         name = type_str
     else:
         raise ValueError(f"array_name_format {_array_name_format} not recognised")
-    metaclass = _make_metaclass(type(array_type))
-    out = metaclass(
-        name,
-        (array_type, AbstractArray),
-        dict(
-            array_type=array_type,
-            dtypes=dtypes,
-            dims=dims,
-            index_variadic=index_variadic,
-            dim_str=dim_str,
-        ),
-    )
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
-        out.__module__ = "builtins"
-    else:
-        out.__module__ = "jaxtyping"
+
+    return (array_type, name, dtypes, dims, index_variadic, dim_str)
+
+
+def _make_array(*args, **kwargs):
+    out = _make_array_cached(*args, **kwargs)
+
+    if type(out) is tuple:
+        array_type, name, dtypes, dims, index_variadic, dim_str = out
+        metaclass = _make_metaclass(type(array_type))
+
+        out = metaclass(
+            name,
+            (array_type, AbstractArray),
+            dict(
+                array_type=array_type,
+                dtypes=dtypes,
+                dims=dims,
+                index_variadic=index_variadic,
+                dim_str=dim_str,
+            ),
+        )
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+
     return out
 
 
@@ -641,10 +674,12 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
 _prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
+_uint4 = "uint4"
 _uint8 = "uint8"
 _uint16 = "uint16"
 _uint32 = "uint32"
 _uint64 = "uint64"
+_int4 = "int4"
 _int8 = "int8"
 _int16 = "int16"
 _int32 = "int32"
@@ -670,10 +705,12 @@ def _make_dtype(_dtypes, name):
     return _Cls
 
 
+UInt4 = _make_dtype(_uint4, "UInt4")
 UInt8 = _make_dtype(_uint8, "UInt8")
 UInt16 = _make_dtype(_uint16, "UInt16")
 UInt32 = _make_dtype(_uint32, "UInt32")
 UInt64 = _make_dtype(_uint64, "UInt64")
+Int4 = _make_dtype(_int4, "Int4")
 Int8 = _make_dtype(_int8, "Int8")
 Int16 = _make_dtype(_int16, "Int16")
 Int32 = _make_dtype(_int32, "Int32")
@@ -686,8 +723,8 @@ Complex64 = _make_dtype(_complex64, "Complex64")
 Complex128 = _make_dtype(_complex128, "Complex128")
 
 bools = [_bool, _bool_]
-uints = [_uint8, _uint16, _uint32, _uint64]
-ints = [_int8, _int16, _int32, _int64]
+uints = [_uint4, _uint8, _uint16, _uint32, _uint64]
+ints = [_int4, _int8, _int16, _int32, _int64]
 floats = [_bfloat16, _float16, _float32, _float64]
 complexes = [_complex64, _complex128]
 
@@ -706,10 +743,4 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num")
 
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
-if has_jax:
-    Key = _make_dtype(_prng_key, "Key")
-    # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
-    # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
-    PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
-    Scalar = Shaped[jax.Array, ""]
-    ScalarLike = Shaped[jax.typing.ArrayLike, ""]
+Key = _make_dtype(_prng_key, "Key")

jaxtyping/_decorator.py

@@ -19,24 +19,18 @@
 
 import dataclasses
 import functools as ft
+import importlib.util
 import inspect
 import itertools as it
 import sys
 import warnings
 from typing import Any, get_args, get_origin, get_type_hints, overload
 
-
-try:
-    import jax._src.traceback_util as traceback_util
-except ImportError:
-    pass
-else:
-    traceback_util.register_exclusion(__file__)
-
+from jaxtyping import AbstractArray
 
 from ._config import config
 from ._errors import AnnotationError, TypeCheckError
-from ._storage import pop_shape_memo, push_shape_memo
+from ._storage import pop_shape_memo, push_shape_memo, shape_str
 
 
 class _Sentinel:
@@ -45,6 +39,7 @@ class _Sentinel:
 
 
 _sentinel = _Sentinel()
+_tb_flag = True
 
 
 @overload
@@ -79,16 +74,29 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
             return x[:, :, None] * y[:, None, :]
 
         # Type-check a dataclass
+        from dataclasses import dataclass
+
         @jaxtyped(typechecker=typechecker)
         @dataclass
         class MyDataclass:
             x: int
-            y: Float[Array "b c"]
+            y: Float[Array, "b c"]
         ```
 
     **Arguments:**
 
-    - `fn`: The function or dataclass to decorate.
+    - `fn`: The function or dataclass to decorate. In practice if you want to use
+        dataclasses with JAX, then
+        [`equinox.Module`](https://docs.kidger.site/equinox/api/module/module/) is our
+        recommended approach:
+        ```python
+        import equinox as eqx
+
+        @jaxtyped(typechecker=typechecker)
+        class MyModule(eqx.Module):
+            ...
+        ```
+
     - `typechecker`: Keyword-only argument: the runtime type-checker to use. This should
         be a function decorator that will raise an exception if there is a type error,
         e.g.
@@ -180,6 +188,18 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
         useful when working at the global scope.
     """
 
+    global _tb_flag
+    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__)
+        _tb_flag = False
+
     # First handle the `jaxtyped("context")` usage, which is a special case.
     if fn == "context":
         if typechecker is not _sentinel:
@@ -296,6 +316,27 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
             # in which case make a best-effort attempt to add shape information for any
             # type errors.
 
+            # we want to detect generators, and ignore return annotations on them,
+            # to avoid issues with O(n) typechecking trying to typecheck yielded values
+            wrp = fn
+            while hasattr(wrp, "__wrapped__"):
+                wrp = wrp.__wrapped__
+
+            if inspect.isgeneratorfunction(wrp) or inspect.isasyncgenfunction(wrp):
+                # recursively parse all the annotations, and mark all the jaxtyping
+                # annotations as not needing instance checks, while still being
+                # visible as original ones for the typechecker
+                def modify_annotation(ann):
+                    if inspect.isclass(ann) and issubclass(ann, AbstractArray):
+                        ann.make_transparent()
+
+                    for sub_ann in get_args(ann):
+                        modify_annotation(sub_ann)
+
+                # just to make sure: check that fn has valid return annotations
+                if hasattr(fn, "__annotations__") and "return" in fn.__annotations__:
+                    modify_annotation(fn.__annotations__["return"])
+
             signature = inspect.signature(fn)
 
             @ft.wraps(fn)
@@ -305,8 +346,9 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
                 try:
                     return fn(*args, **kwargs)
                 except Exception as e:
+                    # add_note api is support from python 3.11+
                     if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
-                        shape_info = _exc_shape_info(memos)
+                        shape_info = shape_str(memos)
                         if shape_info != "":
                             msg = (
                                 "The preceding error occurred within the scope of a "
@@ -398,7 +440,7 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
                             "----------------------\n"
                             f"Called with parameters: {param_values}\n"
                             f"Parameter annotations: {param_hints}.\n"
-                            + _exc_shape_info(memos)
+                            + shape_str(memos)
                         )
                         if config.jaxtyping_remove_typechecker_stack:
                             raise TypeCheckError(msg) from None
@@ -451,7 +493,7 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
                                 "----------------------\n"
                                 f"Called with parameters: {param_values}\n"
                                 f"Parameter annotations: {param_hints}.\n"
-                                + _exc_shape_info(memos)
+                                + shape_str(memos)
                             )
                             if config.jaxtyping_remove_typechecker_stack:
                                 raise TypeCheckError(msg) from None
@@ -476,7 +518,7 @@ class _JaxtypingContext:
 def _check_dataclass_annotations(self, typechecker):
     """Creates and calls a function that checks the attributes of `self`
 
-    `self` should be a dataclass instancae. `typechecker` should be e.g.
+    `self` should be a dataclass instance. `typechecker` should be e.g.
     `beartype.beartype` or `typeguard.typechecked`.
     """
     parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)]
@@ -743,40 +785,6 @@ def _pformat(x, short_self: bool):
     return pformat(x)
 
 
-def _exc_shape_info(memos) -> str:
-    """Gives debug information on the current state of jaxtyping's internal memos.
-    Used in type-checking error messages.
-    """
-    single_memo, variadic_memo, pytree_memo, _ = memos
-    single_memo = {
-        name: size
-        for name, size in single_memo.items()
-        if not name.startswith("~~delete~~")
-    }
-    variadic_memo = {
-        name: shape
-        for name, (_, shape) in variadic_memo.items()
-        if not name.startswith("~~delete~~")
-    }
-    pieces = []
-    if len(single_memo) > 0 or len(variadic_memo) > 0:
-        pieces.append(
-            "The current values for each jaxtyping axis annotation are as follows."
-        )
-        for name, size in single_memo.items():
-            pieces.append(f"{name}={size}")
-        for name, shape in variadic_memo.items():
-            pieces.append(f"{name}={shape}")
-    if len(pytree_memo) > 0:
-        pieces.append(
-            "The current values for each jaxtyping PyTree structure annotation are as "
-            "follows."
-        )
-        for name, structure in pytree_memo.items():
-            pieces.append(f"{name}={structure}")
-    return "\n".join(pieces)
-
-
 class _jaxtyping_note_str(str):
     """Used with `_no_jaxtyping_note` to flag that a note came from jaxtyping."""
 

jaxtyping/_import_hook.py

@@ -358,28 +358,6 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
                                                  # so will be hook'd.
         ```
 
-    !!! warning
-
-        Stringified dataclass annotations, e.g.
-        ```python
-        @dataclass()
-        class Foo:
-            x: "int"
-        ```
-        will be silently skipped without checking them. This is because these are
-        essentially impossible to resolve at runtime. Such stringified annotations
-        typically occur either when using them for forward references, or when using
-        `from __future__ import annotations`. (You should never use the latter, it is
-        largely incompatible with runtime type checking.)
-
-        Partially stringified dataclass annotations, e.g.
-        ```python
-        @dataclass()
-        class Foo:
-            x: tuple["int"]
-        ```
-        will likely raise an error, and must not be used at all.
-
     !!! warning
 
         If a function already has any decorators on it, then `@jaxtyped` will get added

jaxtyping/_indirection.py

@@ -32,6 +32,7 @@ from typing import (
     Annotated as Float64,  # noqa: F401
     Annotated as Inexact,  # noqa: F401
     Annotated as Int,  # noqa: F401
+    Annotated as Int4,  # noqa: F401
     Annotated as Int8,  # noqa: F401
     Annotated as Int16,  # noqa: F401
     Annotated as Int32,  # noqa: F401
@@ -42,12 +43,18 @@ from typing import (
     Annotated as Real,  # noqa: F401
     Annotated as Shaped,  # noqa: F401
     Annotated as UInt,  # noqa: F401
+    Annotated as UInt4,  # noqa: F401
     Annotated as UInt8,  # noqa: F401
     Annotated as UInt16,  # noqa: F401
     Annotated as UInt32,  # noqa: F401
     Annotated as UInt64,  # noqa: F401
+    TYPE_CHECKING,
 )
 
+
+if not TYPE_CHECKING:
+    assert False
+
 from jax import (
     Array as PRNGKeyArray,  # noqa: F401
     Array as Scalar,  # noqa: F401

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)

jaxtyping/_storage.py

@@ -71,6 +71,62 @@ def pop_shape_memo() -> None:
     _shape_storage.memo_stack.pop()
 
 
+def shape_str(memos) -> str:
+    """Gives debug information on the current state of jaxtyping's internal memos.
+    Used in type-checking error messages.
+
+    **Arguments:**
+
+    - `memos`: as returned by `get_shape_memo` or `push_shape_memo`.
+    """
+    single_memo, variadic_memo, pytree_memo, _ = memos
+    single_memo = {
+        name: size
+        for name, size in single_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    variadic_memo = {
+        name: shape
+        for name, (_, shape) in variadic_memo.items()
+        if not name.startswith("~~delete~~")
+    }
+    pieces = []
+    if len(single_memo) > 0 or len(variadic_memo) > 0:
+        pieces.append(
+            "The current values for each jaxtyping axis annotation are as follows."
+        )
+        for name, size in single_memo.items():
+            pieces.append(f"{name}={size}")
+        for name, shape in variadic_memo.items():
+            pieces.append(f"{name}={shape}")
+    if len(pytree_memo) > 0:
+        pieces.append(
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        )
+        for name, structure in pytree_memo.items():
+            pieces.append(f"{name}={structure}")
+    return "\n".join(pieces)
+
+
+def print_bindings():
+    """Prints the values of the current jaxtyping axis bindings. Intended for debugging.
+
+    That is, whilst doing runtime type checking, so that e.g. the `foo` and `bar` of
+    `Float[Array, "foo bar"]` are assigned values -- this function will print out those
+    values.
+
+    **Arguments:**
+
+    Nothing.
+
+    **Returns:**
+
+    Nothing.
+    """
+    print(shape_str(get_shape_memo()))
+
+
 _treepath_storage = threading.local()
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.25"
+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"
@@ -23,7 +23,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 urls = {repository = "https://github.com/google/jaxtyping" }
-dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3,<3", "typing_extensions>=3.7.4.1"]
+dependencies = ["numpy>=1.20.0", "typeguard==2.13.3"]
 entry-points = {pytest11 = {jaxtyping = "jaxtyping._pytest_plugin"}}
 
 [build-system]

test/requirements.txt

@@ -4,5 +4,6 @@ equinox
 IPython
 jaxlib
 pytest
+pytest-asyncio
 tensorflow
 typeguard<3

test/test_all_importable.py

@@ -0,0 +1,47 @@
+# We have some pretty complicated semantics in `__init__.py`.
+# Here we check that we didn't miss one of them on our runtime branch.
+def test_all_importable():
+    # Ordered according to their appearance in the documentation.
+    from jaxtyping import (  # noqa: I001
+        Shaped,  # noqa: F401
+        Bool,  # noqa: F401
+        Key,  # noqa: F401
+        Num,  # noqa: F401
+        Inexact,  # noqa: F401
+        Float,  # noqa: F401
+        BFloat16,  # noqa: F401
+        Float16,  # noqa: F401
+        Float32,  # noqa: F401
+        Float64,  # noqa: F401
+        Complex,  # noqa: F401
+        Complex64,  # noqa: F401
+        Complex128,  # noqa: F401
+        Integer,  # noqa: F401
+        UInt,  # noqa: F401
+        UInt4,  # noqa: F401
+        UInt8,  # noqa: F401
+        UInt16,  # noqa: F401
+        UInt32,  # noqa: F401
+        UInt64,  # noqa: F401
+        Int,  # noqa: F401
+        Int4,  # noqa: F401
+        Int8,  # noqa: F401
+        Int16,  # noqa: F401
+        Int32,  # noqa: F401
+        Int64,  # noqa: F401
+        Real,  # noqa: F401
+        Array,  # noqa: F401
+        ArrayLike,  # noqa: F401
+        Scalar,  # noqa: F401
+        ScalarLike,  # noqa: F401
+        PRNGKeyArray,  # noqa: F401
+        PyTreeDef,  # noqa: F401
+        PyTree,  # noqa: F401
+        jaxtyped,  # noqa: F401
+        install_import_hook,  # noqa: F401
+        AbstractArray,  # noqa: F401
+        AbstractDtype,  # noqa: F401
+        print_bindings,  # noqa: F401
+        get_array_name_format,  # noqa: F401
+        set_array_name_format,  # noqa: F401
+    )

test/test_array.py

@@ -25,7 +25,12 @@ import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import pytest
-import torch
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
 
 from jaxtyping import (
     AbstractDtype,
@@ -67,6 +72,7 @@ def test_dtypes():
         Float64,
         Inexact,
         Int,
+        Int4,
         Int8,
         Int16,
         Int32,
@@ -74,6 +80,7 @@ def test_dtypes():
         Num,
         Shaped,
         UInt,
+        UInt4,
         UInt8,
         UInt16,
         UInt32,
@@ -125,7 +132,9 @@ def test_any_dtype(jaxtyp, typecheck, getkey):
 
     g(jr.normal(getkey(), (3, 4)))
     g(jnp.array([[True, False]]))
+    g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int4))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int8))
+    g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint4))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint16))
     g(jr.normal(getkey(), (3, 4), dtype=jnp.complex128))
     g(jr.normal(getkey(), (3, 4), dtype=jnp.bfloat16))
@@ -549,7 +558,9 @@ def test_arraylike(typecheck, getkey):
 def test_subclass():
     assert issubclass(Float[Array, ""], Array)
     assert issubclass(Float[np.ndarray, ""], np.ndarray)
-    assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
+
+    if torch is not None:
+        assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
 
 
 def test_ignored_names():

test/test_decorator.py

@@ -1,9 +1,10 @@
 import abc
 
+import jax.numpy as jnp
 import jax.random as jr
 import pytest
 
-from jaxtyping import Array, Float, jaxtyped
+from jaxtyping import Array, Float, jaxtyped, print_bindings
 
 from .helpers import ParamError, ReturnError
 
@@ -166,3 +167,17 @@ def test_local_stringified_annotation(typecheck):
     # We don't check that errors are raised if it goes wrong, since we can't usually
     # resolve local type annotations at runtime. Best we can hope for is not to raise
     # a spurious error about not being able to find the type.
+
+
+def test_print_bindings(typecheck, capfd):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: Float[Array, "foo bar"]):
+        print_bindings()
+
+    capfd.readouterr()
+    f(jnp.zeros((3, 4)))
+    text, _ = capfd.readouterr()
+    assert text == (
+        "The current values for each jaxtyping axis annotation are as follows."
+        "\nfoo=3\nbar=4\n"
+    )

test/test_equals.py

@@ -0,0 +1,35 @@
+from typing import Tuple, Union
+
+import pytest
+
+from jaxtyping import (
+    Array,
+    Float,
+    Float32,
+    Integer,
+    PRNGKeyArray,
+    PyTree,
+    Shaped,
+)
+
+
+@pytest.mark.parametrize(
+    "make_fn",
+    [
+        lambda: Float[Array, "4"],
+        lambda: Float32[Array, ""],
+        lambda: Integer[Array, "1 2 3"],
+        lambda: Shaped[PRNGKeyArray, "2"],
+        lambda: Float[float, "#*shape"],
+        lambda: PyTree[int],
+        lambda: PyTree[Float[Array, ""]],
+        lambda: PyTree[Float32[Array, "*m b c"]],
+        lambda: PyTree[PyTree[Float32[Array, "1 2 b *"]]],
+        lambda: PyTree[Union[str, Float32[Array, "1"]]],
+        lambda: PyTree[
+            Tuple[int, float, Float[Array, ""], PyTree[Union[Float[Array, ""], float]]]
+        ],
+    ],
+)
+def test_equals(make_fn):
+    assert make_fn() == make_fn()

test/test_generators.py

@@ -0,0 +1,88 @@
+from typing import AsyncIterator, Iterator
+
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Array, Float, Shaped
+
+from .helpers import ParamError
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+def test_generators_simple(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def gen(x: Float[Array, "*"]) -> Iterator[Float[Array, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def foo():
+        next(gen(jnp.zeros(2)))
+        next(gen(jnp.zeros((3, 4))))
+
+    foo()
+
+
+def test_generators_return_no_annotations(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def gen(x: Float[Array, "*"]):
+        yield x
+
+    @jaxtyp(typecheck)
+    def foo():
+        next(gen(jnp.zeros(2)))
+        next(gen(jnp.zeros((3, 4))))
+
+    foo()
+
+
+@pytest.mark.asyncio
+async def test_async_generators_simple(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    async def gen(x: Float[Array, "*"]) -> AsyncIterator[Float[Array, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    async def foo():
+        async for _ in gen(jnp.zeros(2)):
+            pass
+        async for _ in gen(jnp.zeros((3, 4))):
+            pass
+
+    await foo()
+
+
+def test_generators_dont_modify_same_annotations(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def g(x: Float[Array, "1"]) -> Iterator[Float[Array, "1"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def m(x: Float[Array, "1"]) -> Float[Array, "1"]:
+        return x
+
+    with pytest.raises(ParamError):
+        next(g(jnp.zeros(2)))
+    with pytest.raises(ParamError):
+        m(jnp.zeros(2))
+
+
+def test_generators_original_issue(jaxtyp, typecheck):
+    # Effectively the same as https://github.com/patrick-kidger/jaxtyping/issues/91
+    if torch is None:
+        pytest.skip("torch is not available")
+
+    @jaxtyp(typecheck)
+    def g(x: Shaped[torch.Tensor, "*"]) -> Iterator[Shaped[torch.Tensor, "*"]]:
+        yield x
+
+    @jaxtyp(typecheck)
+    def f():
+        next(g(torch.zeros(1)))
+        next(g(torch.zeros(2)))
+
+    f()

test/test_no_jax_dependency.py

@@ -0,0 +1,25 @@
+import subprocess
+import sys
+
+
+_py_path = sys.executable
+
+
+def test_no_jax_dependency():
+    result = subprocess.run(
+        f"{_py_path} -c "
+        "'import jaxtyping; import sys; sys.exit(\"jax\" in sys.modules)'",
+        shell=True,
+    )
+    assert result.returncode == 0
+
+
+# Meta-test: test that the above test will work. (i.e. that I haven't messed up using
+# subprocess.)
+def test_meta():
+    result = subprocess.run(
+        f"{_py_path} -c 'import jaxtyping; import jax; import sys; "
+        'sys.exit("jax" in sys.modules)\'',
+        shell=True,
+    )
+    assert result.returncode == 1

test/test_serialisation.py

@@ -1,16 +1,25 @@
 import cloudpickle
 import numpy as np
-import torch
+
+
+try:
+    import torch
+except ImportError:
+    torch = None
 
 from jaxtyping import AbstractArray, Array, Shaped
 
 
 def test_pickle():
     x = cloudpickle.dumps(Shaped[Array, ""])
-    y = cloudpickle.dumps(AbstractArray)
-    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
-    w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
     cloudpickle.loads(x)
+
+    y = cloudpickle.dumps(AbstractArray)
     cloudpickle.loads(y)
+
+    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
     cloudpickle.loads(z)
-    cloudpickle.loads(w)
+
+    if torch is not None:
+        w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
+        cloudpickle.loads(w)