link to why use type annotations

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

.pre-commit-config.yaml

@@ -18,12 +18,11 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 repos:
-  - repo: https://github.com/ambv/black
+  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 23.9.1
+    rev: v0.1.7
     hooks:
-    - id: black
+      - id: ruff  # linter
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
+        types_or: [ python, pyi, jupyter ]
-    rev: 'v0.0.291'
+        args: [ --fix ]
-    hooks:
+      - id: ruff-format  # formatter
-      - id: ruff
+        types_or: [ python, pyi, jupyter ]
-        args: ["--fix"]

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/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!  
 
+**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.  

docs/requirements.txt

@@ -3,7 +3,7 @@ mkdocs-material==7.3.6   # Theme
 pymdown-extensions==9.4  # Markdown extensions e.g. to handle LaTeX.
 mkdocstrings==0.17.0     # Autogenerate documentation from docstrings.
 mknotebooks==0.7.1       # Turn Jupyter Lab notebooks into webpages.
-pytkdocs_tweaks==0.0.5   # Tweaks mkdocstrings to improve various aspects
+pytkdocs_tweaks==0.0.8   # Tweaks mkdocstrings to improve various aspects
 mkdocs_include_exclude_files==0.0.1  # Tweak which files are included/excluded
 jinja2==3.0.3            # Older version. After 3.1.0 seems to be incompatible with current versions of mkdocstrings.
 pygments==2.14.0

jaxtyping/__init__.py

@@ -17,56 +17,39 @@
 # 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 ._decorator import jaxtyped as jaxtyped, TypeCheckError as TypeCheckError
+from ._config import config as config
+from ._decorator import jaxtyped as jaxtyped
+from ._errors import (
+    AnnotationError as AnnotationError,
+    TypeCheckError as TypeCheckError,
+)
 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
+    import typing_extensions
-    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+
     from jax import Array as Array
-    from jax.typing import ArrayLike as ArrayLike
+    from jax.tree_util import PyTreeDef as PyTreeDef
-elif has_jax:
+    from jax.typing import ArrayLike as ArrayLike, DTypeLike as DTypeLike
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
 
-        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,
@@ -79,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,
@@ -86,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")
@@ -155,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:
+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,
+    )
+
+    # 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:
+
+        @ft.cache
+        def __getattr__(item):
+            if item == "Array":
+                if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+                    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
+                    # `typing.GENERATING_DOCUMENTATION = True` when generating
-        # add whatever shims are necessary to get pretty docs. E.g. to have type
+                    # documentation, to add whatever shims are necessary to get pretty
-        # annotations appear as just `PyTree`, not `jaxtyping.PyTree`.
+                    # 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`,
+                    # As jaxtyping actually wants things to appear as e.g.
-        # rather than just `PyTree`, then it sets
+                    # `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
+                    # Here we do only a `hasattr` check, as we want to get this version
-        # `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs.
+                    # 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
+                        """Alias for `jax.tree_util.PyTreeDef`, which is the type of the
-            from `jax.tree_util.tree_structure(...)`.
+                        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:
-        from jax.tree_util import PyTreeDef as PyTreeDef
+                    import jax.tree_util
 
-    from ._pytree_type import PyTree as PyTree  # noqa: F401
+                    return jax.tree_util.PyTreeDef
-
-
-# Conveniences
-if typing.TYPE_CHECKING:
-    from ._indirection import (
-        PRNGKeyArray as PRNGKeyArray,
-        Scalar as Scalar,
-        ScalarLike as ScalarLike,
-    )
-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
-
-        PRNGKeyArray.__module__ = "builtins"
             else:
-        from ._array_types import PRNGKeyArray
+                raise AttributeError(f"module jaxtyping has no attribute {item!r}")
-
-del has_jax
 
 
 check_equinox_version = True  # easy-to-replace line with copybara

jaxtyping/_array_types.py

@@ -23,11 +23,12 @@ import re
 import sys
 import types
 import typing
+from dataclasses import dataclass
 from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
-from ._raise import jaxtyping_raise, jaxtyping_raise_from
+from ._errors import AnnotationError
 from ._storage import (
     get_shape_memo,
     get_treeflatten_memo,
@@ -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):
+    name: str
-        self.name = name
+    broadcastable: bool
-        self.broadcastable = broadcastable
+    treepath: Any
-        self.treepath = treepath
 
 
+@dataclass(frozen=True)
 class _NamedVariadicDim:
-    def __init__(self, name, broadcastable, treepath):
+    name: str
-        self.name = name
+    broadcastable: bool
-        self.broadcastable = broadcastable
+    treepath: Any
-        self.treepath = treepath
 
 
+@dataclass(frozen=True)
 class _FixedDim:
-    def __init__(self, size, broadcastable):
+    size: str
-        self.size = size
+    broadcastable: bool
-        self.broadcastable = broadcastable
 
 
+@dataclass(frozen=True)
 class _SymbolicDim:
-    def __init__(self, elem, broadcastable):
+    elem: Any
-        self.elem = elem
+    broadcastable: bool
-        self.broadcastable = broadcastable
 
 
 _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.
@@ -133,18 +121,15 @@ def _check_dims(
                 # Make a copy to avoid `__builtins__` getting added as a key.
                 eval_size = eval(elem, single_memo.copy())
             except NameError as e:
-                jaxtyping_raise_from(
+                raise AnnotationError(
-                    NameError(
                     f"Cannot process symbolic axis '{cls_dim.elem}' as "
                     "some axis names have not been processed. In practice you "
                     "should usually only use symbolic axes in annotations "
                     "for return types, referring only to axes annotated for "
                     "arguments."
-                    ),
+                ) from e
-                    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:
@@ -157,16 +142,26 @@ def _check_dims(
                 single_memo[name] = obj_size
             else:
                 if cls_size != obj_size:
-                    return False
+                    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 True
+    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
@@ -180,8 +175,8 @@ class _MetaAbstractArray(type):
             if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
                 dtype = repr_dtype[1]
             else:
-                jaxtyping_raise(
+                raise AnnotationError(
-                    RuntimeError("Unrecognised array/tensor type to extract dtype from")
+                    "Unrecognised array/tensor type to extract dtype from"
                 )
 
         if cls.dtypes is not _any_dtype:
@@ -196,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()
@@ -210,13 +208,13 @@ class _MetaAbstractArray(type):
                 single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
             )
             raise
-        if check:
+        if check == "":
-            return True
+            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,
@@ -224,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):
+            if len(obj.shape) != len(cls.dims):
-                return False
+                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 False
+                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):
+            prefix_check = _check_dims(
-                return False
+                cls.dims[:i], obj.shape[:i], single_memo, arg_memo
-            if j is not None and not _check_dims(
+            )
+            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
-            ):
+                )
-                return False
+                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:
@@ -256,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:
@@ -274,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 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 True
+                    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
 
@@ -318,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 "
@@ -540,7 +559,17 @@ def _make_array(array_type, dim_str, dtypes, name):
         name = type_str
     else:
         raise ValueError(f"array_name_format {_array_name_format} not recognised")
+
+    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),
@@ -556,18 +585,17 @@ def _make_array(array_type, dim_str, dtypes, name):
             out.__module__ = "builtins"
         else:
             out.__module__ = "jaxtyping"
+
     return out
 
 
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
-        jaxtyping_raise(
+        raise AnnotationError(
-            RuntimeError(
             f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
             "check just the dtype of an array, then use "
             f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
         )
-        )
 
     def __getitem__(cls, item: tuple[Any, str]):
         if not isinstance(item, tuple) or len(item) != 2:
@@ -646,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"
@@ -675,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")
@@ -691,8 +723,8 @@ Complex64 = _make_dtype(_complex64, "Complex64")
 Complex128 = _make_dtype(_complex128, "Complex128")
 
 bools = [_bool, _bool_]
-uints = [_uint8, _uint16, _uint32, _uint64]
+uints = [_uint4, _uint8, _uint16, _uint32, _uint64]
-ints = [_int8, _int16, _int32, _int64]
+ints = [_int4, _int8, _int16, _int32, _int64]
 floats = [_bfloat16, _float16, _float32, _float64]
 complexes = [_complex64, _complex128]
 
@@ -711,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, ""]

jaxtyping/_config.py

@@ -0,0 +1,48 @@
+import os
+from typing import Union
+
+
+def _maybestr2bool(value: Union[bool, str], error: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    elif isinstance(value, str):
+        if value.lower() in ("0", "false"):
+            return False
+        elif value.lower() in ("1", "true"):
+            return True
+        else:
+            raise ValueError(error)
+    else:
+        raise ValueError(error)
+
+
+class _JaxtypingConfig:
+    def __init__(self):
+        self.update("jaxtyping_disable", os.environ.get("JAXTYPING_DISABLE", "0"))
+        self.update(
+            "jaxtyping_remove_typechecker_stack",
+            os.environ.get("JAXTYPING_REMOVE_TYPECHECKER_STACK", "0"),
+        )
+
+    def update(self, item: str, value):
+        if item.lower() == "jaxtyping_disable":
+            msg = (
+                "Unrecognised value for `JAXTYPING_DISABLE`. Valid values are "
+                "`JAXTYPING_DISABLE=0` (the default) or `JAXTYPING_DISABLE=1` (to "
+                "disable runtime type checking)."
+            )
+            self.jaxtyping_disable = _maybestr2bool(value, msg)
+        elif item.lower() == "jaxtyping_remove_typechecker_stack":
+            msg = (
+                "Unrecognised value for `JAXTYPING_REMOVE_TYPECHECKER_STACK`. Valid "
+                "values are `JAXTYPING_REMOVE_TYPECHECKER_STACK=0` (the default) or "
+                "`JAXTYPING_REMOVE_TYPECHECKER_STACK=1` (to remove the stack frames "
+                "from the typechecker in `jaxtyped(typechecker=...)`, when it raises a "
+                "runtime type-checking error)."
+            )
+            self.jaxtyping_remove_typechecker_stack = _maybestr2bool(value, msg)
+        else:
+            raise ValueError(f"Unrecognised config value {item}")
+
+
+config = _JaxtypingConfig()

jaxtyping/_decorator.py

@@ -19,46 +19,40 @@
 
 import dataclasses
 import functools as ft
+import importlib.util
 import inspect
 import itertools as it
 import sys
-import types
+import warnings
-import weakref
 from typing import Any, get_args, get_origin, get_type_hints, overload
 
+from jaxtyping import AbstractArray
 
-try:
+from ._config import config
-    import jax._src.traceback_util as traceback_util
+from ._errors import AnnotationError, TypeCheckError
-except ImportError:
+from ._storage import pop_shape_memo, push_shape_memo, shape_str
-    pass
-else:
-    traceback_util.register_exclusion(__file__)
 
 
-from ._storage import pop_shape_memo, push_shape_memo
+class _Sentinel:
+    def __repr__(self):
+        return "sentinel"
 
 
-_jaxtyped_fns = weakref.WeakSet()
+_sentinel = _Sentinel()
-
+_tb_flag = True
-
-class TypeCheckError(TypeError):
-    pass
-
-
-TypeCheckError.__module__ = "jaxtyping"  # appears in error messages
 
 
 @overload
-def jaxtyped(*, typechecker=None):
+def jaxtyped(*, typechecker=_sentinel):
     ...
 
 
 @overload
-def jaxtyped(fn, *, typechecker=None):
+def jaxtyped(fn, *, typechecker=_sentinel):
     ...
 
 
-def jaxtyped(fn=None, *, typechecker=None):
+def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
     """Decorate a function with this to perform runtime type-checking of its arguments
     and return value. Decorate a dataclass to perform type-checking of its attributes.
 
@@ -80,18 +74,32 @@ def jaxtyped(fn=None, *, typechecker=None):
             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
-    - `typechecker`: The runtime type-checker to use. This should be a function
+        dataclasses with JAX, then
-        decorator that will raise an exception if there is a type error, e.g.
+        [`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.
         ```python
         @typechecker
         def f(x: int):
@@ -104,7 +112,7 @@ def jaxtyped(fn=None, *, typechecker=None):
         skip automatic runtime type-checking, but still support manual `isinstance`
         checks inside the function body:
         ```python
-        @jaxtyped
+        @jaxtyped(typechecker=None)
         def f(x):
             assert isinstance(x, Float[Array, "batch channel"])
         ```
@@ -126,10 +134,10 @@ def jaxtyped(fn=None, *, typechecker=None):
         @typechecker
         def f(...): ...
         ```
-        This is still supported, but the `jaxtyped(typechecker=typechecker)` syntax
+        This is still supported, but will now raise a warning recommending the
-        discussed above will produce easier-to-debug error messages. Under the hood, the
+        `jaxtyped(typechecker=typechecker)` syntax discussed above. (Which will produce
-        new syntax more carefully manipulates the typechecker so as to determine where
+        easier-to-debug error messages: under the hood, the new syntax more carefully
-        a type-check error arises.
+        manipulates the typechecker so as to determine where a type-check error arises.)
 
     ??? Info "Notes for advanced users"
 
@@ -163,26 +171,88 @@ def jaxtyped(fn=None, *, typechecker=None):
 
         **Decoupling contexts from function calls:**
 
-        If you would like a new dynamic context *without* calling a new function, then
+        If you would like to call a new function *without* creating a new
-        `jaxtyped` may be passed the string `"context"` and used as a context manager:
+        dynamic context (and using the same set of axis and structure values), then
+        simply do not add a `jaxtyped` decorator to your inner function, whilst
+        continuing to perform type-checking in whatever way you prefer.
+
+        Conversely, if you would like a new dynamic context *without* calling a new
+        function, then in addition to the usage discussed above, `jaxtyped` also
+        supports being used as a context manager, by passing it the string `"context"`:
         ```python
         with jaxtyped("context"):
             assert isinstance(x, Float[Array, "batch channel"])
         ```
-        which is equivalent to placing this code inside a new function wrapped in
+        This is equivalent to placing this code inside a new function wrapped in
         `jaxtyped(typechecker=None)`. Usage like this is very rare; it's mostly only
         useful when working at the global scope.
-
-        Conversely, if you would like to call a new function *without* creating a new
-        dynamic context (and using the same set of axis and structure values), then
-        simply do not add a `jaxtyped` decorator to your inner function, whilst
-        continuing to perform type-checking in whatever way you prefer.
     """
 
-    if fn is None:
+    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:
+            raise ValueError(
+                "Cannot use `jaxtyped` as a context with a typechecker. That is, "
+                "`with jaxtyped('context', typechecker=...):`. is not allowed. In this "
+                "case the type checker does not actually do anything, as there is no "
+                "function to type-check."
+            )
+        return _JaxtypingContext()
+
+    # Now check that a typechecker has been explicitly declared. (Or explicitly declared
+    # as not being used, via `typechecker=None`.)
+    # This is needed just for backward compatibility: an undeclared typechecker
+    # corresponds to the old double-decorator syntax.
+    if typechecker is _sentinel:
+        # This branch will also catch the easy-to-make mistake of
+        # ```python
+        # @jaxtyped(typechecker)
+        # def foo(...):
+        # ```
+        # which is a bug as `typechecker` is interpreted as the function to decorate!
+        warnings.warn(
+            "As of jaxtyping version 0.2.24, jaxtyping now prefers the syntax\n"
+            "```\n"
+            "from jaxtyping import jaxtyped\n"
+            "# Use your favourite typechecker: usually one of the two lines below.\n"
+            "from typeguard import typechecked as typechecker\n"
+            "from beartype import beartype as typechecker\n"
+            "\n"
+            "@jaxtyped(typechecker=typechecker)\n"
+            "def foo(...):\n"
+            "```\n"
+            "and the old double-decorator syntax\n"
+            "```\n"
+            "@jaxtyped\n"
+            "@typechecker\n"
+            "def foo(...):\n"
+            "```\n"
+            "should no longer be used. (It will continue to work as it did before, but "
+            "the new approach will produce more readable error messages.)\n"
+            "In particular note that `typechecker` must be passed via keyword "
+            "argument; the following is not valid:\n"
+            "```\n"
+            "@jaxtyped(typechecker)\n"
+            "def foo(...):\n"
+            "```\n",
+            stacklevel=2,
+        )
+        typechecker = None
+
+    if fn is _sentinel:
         return ft.partial(jaxtyped, typechecker=typechecker)
-    elif type(fn) is types.FunctionType and fn in _jaxtyped_fns:
-        return fn
     elif inspect.isclass(fn):
         if dataclasses.is_dataclass(fn) and typechecker is not None:
             # This does not check that the arguments passed to `__init__` match the
@@ -235,15 +305,6 @@ def jaxtyped(fn=None, *, typechecker=None):
         else:
             fdel = jaxtyped(fn.fdel, typechecker=typechecker)
         return property(fget=fget, fset=fset, fdel=fdel)
-    elif fn == "context":
-        if typechecker is not None:
-            raise ValueError(
-                "Cannot use `jaxtyped` as a context with a typechecker. That is, "
-                "`with jaxtyped('context', typechecker=...):`. is not allowed. In this "
-                "case the type checker does not actually do anything, as there is no "
-                "function to type-check."
-            )
-        return _JaxtypingContext()
     else:
         if typechecker is None:
             # Probably being used in the old style as
@@ -255,6 +316,27 @@ def jaxtyped(fn=None, *, typechecker=None):
             # 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)
@@ -264,8 +346,9 @@ def jaxtyped(fn=None, *, typechecker=None):
                 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 "
@@ -321,6 +404,9 @@ def jaxtyped(fn=None, *, typechecker=None):
 
             @ft.wraps(fn)
             def wrapped_fn(*args, **kwargs):
+                if config.jaxtyping_disable:
+                    return fn(*args, **kwargs)
+
                 # Raise bind-time errors before we do any shape analysis. (I.e. skip
                 # the pointless jaxtyping information for a non-typechecking failure.)
                 bound = param_signature.bind(*args, **kwargs)
@@ -331,12 +417,16 @@ def jaxtyped(fn=None, *, typechecker=None):
                     # called.
                     try:
                         param_fn(*args, **kwargs)
-                    except Exception as e:
+                    except AnnotationError:
-                        if hasattr(e, "_jaxtyping_malformed"):
                         raise
-                        else:
+                    except Exception as e:
                         argmsg = _get_problem_arg(
-                                param_signature, args, kwargs, module, typechecker
+                            param_signature,
+                            args,
+                            kwargs,
+                            bound.arguments,
+                            module,
+                            typechecker,
                         )
                         try:
                             name = fn.__name__
@@ -347,10 +437,14 @@ def jaxtyped(fn=None, *, typechecker=None):
                         msg = (
                             "Type-check error whilst checking the parameters of "
                             f"{name}.{argmsg}\n"
-                                f"Called with arguments: {param_values}\n"
+                            "----------------------\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
+                        else:
                             raise TypeCheckError(msg) from e
 
                     # Actually call the function.
@@ -374,17 +468,14 @@ def jaxtyped(fn=None, *, typechecker=None):
                         kwargs[output_name] = out
                         try:
                             full_fn(*args, **kwargs)
-                        except Exception as e:
+                        except AnnotationError:
-                            if hasattr(e, "_jaxtyping_malformed"):
                             raise
-                            else:
+                        except Exception as e:
                             try:
                                 name = fn.__name__
                             except AttributeError:
                                 name = fn.__class__.__name__
-                                param_values = _pformat(
+                            param_values = _pformat(bound.arguments, short_self=True)
-                                    bound.arguments, short_self=True
-                                )
                             return_value = _pformat(out, short_self=False)
                             param_hints = _remove_typing(param_signature)
                             return_hint = _remove_typing(
@@ -397,19 +488,22 @@ def jaxtyped(fn=None, *, typechecker=None):
                             msg = (
                                 "Type-check error whilst checking the return value "
                                 f"of {name}.\n"
-                                    f"Called with arguments: {param_values}\n"
+                                f"Actual value: {return_value}\n"
-                                    f"Return value: {return_value}\n"
+                                f"Expected type: {return_hint}.\n"
+                                "----------------------\n"
+                                f"Called with parameters: {param_values}\n"
                                 f"Parameter annotations: {param_hints}.\n"
-                                    f"Return annotation: {return_hint}.\n"
+                                + shape_str(memos)
-                                    + _exc_shape_info(memos)
                             )
+                            if config.jaxtyping_remove_typechecker_stack:
+                                raise TypeCheckError(msg) from None
+                            else:
                                 raise TypeCheckError(msg) from e
 
                     return out
                 finally:
                     pop_shape_memo()
 
-        _jaxtyped_fns.add(wrapped_fn)
         return wrapped_fn
 
 
@@ -424,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)]
@@ -615,7 +709,7 @@ def _make_argpiece(p, name_to_annotation, name_to_default):
 
 
 def _get_problem_arg(
-    param_signature: inspect.Signature, args, kwargs, module, typechecker
+    param_signature: inspect.Signature, args, kwargs, arguments, module, typechecker
 ) -> str:
     """Determines which argument was likely to be the problematic one responsible for
     raising a type-check error.
@@ -624,13 +718,17 @@ def _get_problem_arg(
     # anyway.
     for keep_name in param_signature.parameters.keys():
         new_parameters = []
+        keep_annotation = sentinel = object()
         for p_name, p in param_signature.parameters.items():
             if p_name == keep_name:
                 new_parameters.append(
                     inspect.Parameter(p.name, p.kind, annotation=p.annotation)
                 )
+                assert keep_annotation is sentinel
+                keep_annotation = _remove_typing(p.annotation)
             else:
                 new_parameters.append(inspect.Parameter(p.name, p.kind))
+        assert keep_annotation is not sentinel
         new_signature = inspect.Signature(new_parameters)
         fn = _make_fn_with_signature(
             "check_single_arg", new_signature, module, output=False
@@ -639,7 +737,12 @@ def _get_problem_arg(
         try:
             fn(*args, **kwargs)
         except Exception:
-            return f"\nThe problem arose whilst typechecking argument '{keep_name}'."
+            keep_value = _pformat(arguments[keep_name], short_self=False)
+            return (
+                f"\nThe problem arose whilst typechecking parameter '{keep_name}'.\n"
+                f"Actual value: {keep_value}\n"
+                f"Expected type: {keep_annotation}."
+            )
     else:
         # Could not localise the problem to a single argument -- probably due to
         # e.g. a mismatched typevar, which each individual argument is okay with.
@@ -682,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/_errors.py

@@ -0,0 +1,12 @@
+class TypeCheckError(TypeError):
+    pass
+
+
+# Not inheriting from TypeError as that gets caught and re-reraised as just a TypeError
+# when using typeguard<3.
+class AnnotationError(Exception):
+    pass
+
+
+TypeCheckError.__module__ = "jaxtyping"
+AnnotationError.__module__ = "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:
+    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:
+        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)
-    except NameError:
-        raise NameError(
-            "ChooseTypecheckerMagics is not defined.\n\n"
-            + "You may be trying to use IPython extension without IPython installed."
-        )

jaxtyping/_pytree_type.py

@@ -24,7 +24,7 @@ from typing import Any, Generic, TypeVar
 import jax.tree_util as jtu
 import typeguard
 
-from ._raise import jaxtyping_raise_from
+from ._errors import AnnotationError
 from ._storage import (
     clear_treeflatten_memo,
     clear_treepath_memo,
@@ -141,14 +141,11 @@ class _MetaPyTree(type):
                     try:
                         prev_structure = pytree_memo[identifier]
                     except KeyError as e:
-                        jaxtyping_raise_from(
+                        raise AnnotationError(
-                            NameError(
                             f"Cannot process composite structure '{cls.structure}' "
                             f"as the structure name {identifier} has not been seen "
                             "before."
-                            ),
+                        ) from e
-                            e,
-                        )
                     # Not using `PyTreeDef.compose` due to JAX bug #18218.
                     prev_pytree = jtu.tree_unflatten(
                         prev_structure, [0] * prev_structure.num_leaves

jaxtyping/_raise.py

@@ -1,23 +0,0 @@
-from typing import NoReturn
-
-
-def jaxtyping_raise(e) -> NoReturn:
-    """Raises `e`, whilst adding a tag that it should not be intercepted by
-    `TypeCheckError`. All `raise` statements from within `__instancecheck__` should use
-    this.
-    """
-    __tracebackhide__ = True
-    try:
-        raise e
-    except Exception as f:
-        f._jaxtyping_malformed = True
-        raise
-
-
-def jaxtyping_raise_from(e, e_base) -> NoReturn:
-    __tracebackhide__ = True
-    try:
-        raise e from e_base
-    except Exception as f:
-        f._jaxtyping_malformed = True
-        raise

jaxtyping/_storage.py

@@ -20,7 +20,7 @@
 import threading
 from typing import Any, Optional
 
-from ._raise import jaxtyping_raise
+from ._errors import AnnotationError
 
 
 _shape_storage = threading.local()
@@ -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()
 
 
@@ -80,13 +136,11 @@ def clear_treepath_memo() -> None:
 
 def set_treepath_memo(index: Optional[int], structure: str) -> None:
     if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
-        jaxtyping_raise(
+        raise AnnotationError(
-            ValueError(
             "Cannot typecheck annotations of the form "
             "`PyTree[PyTree[Shaped[Array, '?foo'], 'T'], 'S']` as it is ambiguous "
             "which PyTree the `?` annotation refers to."
         )
-        )
     if index is None:
         _treepath_storage.value = f"~~delete~~({structure}) "
     else:
@@ -96,13 +150,11 @@ def set_treepath_memo(index: Optional[int], structure: str) -> None:
 
 def get_treepath_memo() -> str:
     if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None:
-        jaxtyping_raise(
+        raise AnnotationError(
-            ValueError(
             "Cannot use `?` annotations, e.g. `Shaped[Array, '?foo']`, except "
             "when contained with structured `PyTree` annotations, e.g. "
             "`PyTree[Shaped[Array, '?foo'], 'T']`."
         )
-        )
     return _treepath_storage.value
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.24"
+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/conftest.py

@@ -57,6 +57,7 @@ def jaxtyp(request):
         # def f(...)
         def impl(typechecker):
             def decorator(fn):
+                with pytest.warns(match="As of jaxtyping version 0.2.24"):
                     return jaxtyping.jaxtyped(typechecker(fn))
 
             return decorator

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,10 +25,16 @@ import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import pytest
+
+
+try:
     import torch
+except ImportError:
+    torch = None
 
 from jaxtyping import (
     AbstractDtype,
+    AnnotationError,
     Array,
     ArrayLike,
     Bool,
@@ -66,6 +72,7 @@ def test_dtypes():
         Float64,
         Inexact,
         Int,
+        Int4,
         Int8,
         Int16,
         Int32,
@@ -73,6 +80,7 @@ def test_dtypes():
         Num,
         Shaped,
         UInt,
+        UInt4,
         UInt8,
         UInt16,
         UInt32,
@@ -124,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))
@@ -448,7 +458,7 @@ def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
         pass
 
     x = jr.normal(getkey(), (4,))
-    with pytest.raises(NameError):
+    with pytest.raises(AnnotationError):
         foo(x)
 
 
@@ -548,6 +558,8 @@ def test_arraylike(typecheck, getkey):
 def test_subclass():
     assert issubclass(Float[Array, ""], Array)
     assert issubclass(Float[np.ndarray, ""], np.ndarray)
+
+    if torch is not None:
         assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
 
 

test/test_decorator.py

@@ -1,57 +1,58 @@
 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
 
 
 class M(metaclass=abc.ABCMeta):
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def f(self):
         ...
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @classmethod
     def g1(cls):
         return 3
 
     @classmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def g2(cls):
         return 4
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @staticmethod
     def h1():
         return 3
 
     @staticmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def h2():
         return 4
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @abc.abstractmethod
     def i1(self):
         ...
 
     @abc.abstractmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def i2(self):
         ...
 
 
 class N:
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @property
     def j1(self):
         return 3
 
     @property
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def j2(self):
         return 4
 
@@ -154,6 +155,8 @@ def test_local_stringified_annotation(typecheck):
 
     f(LocalFoo())
 
+    with pytest.warns(match="As of jaxtyping version 0.2.24"):
+
         @jaxtyped
         @typecheck
         def g(x: "LocalFoo") -> "LocalFoo":
@@ -164,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_messages.py

@@ -14,8 +14,8 @@ def test_arg_localisation(typecheck):
 
     matches = [
         "Type-check error whilst checking the parameters of f",
-        "The problem arose whilst typechecking argument 'z'.",
+        "The problem arose whilst typechecking parameter 'z'.",
-        "Called with arguments: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
+        "Called with parameters: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
         r"Parameter annotations: \(x: str, y: str, z: int\).",
     ]
     for match in matches:
@@ -30,8 +30,8 @@ def test_arg_localisation(typecheck):
     y = jnp.zeros((4, 3))
     matches = [
         "Type-check error whilst checking the parameters of g",
-        "The problem arose whilst typechecking argument 'y'.",
+        "The problem arose whilst typechecking parameter 'y'.",
-        r"Called with arguments: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
+        r"Called with parameters: {'x': f32\[2,3\], 'y': f32\[4,3\]}",
         (
             r"Parameter annotations: \(x: Float\[Array, 'a b'\], y: "
             r"Float\[Array, 'b c'\]\)."
@@ -54,9 +54,9 @@ def test_return(typecheck):
     y = {"a": 1}
     matches = [
         "Type-check error whilst checking the return value of f",
-        r"Called with arguments: {'x': \(1, 2\), 'y': {'a': 1}}",
+        r"Called with parameters: {'x': \(1, 2\), 'y': {'a': 1}}",
-        "Return value: 'foo'",
+        "Actual value: 'foo'",
-        r"Return annotation: PyTree\[Any, \"T S\"\].",
+        r"Expected type: PyTree\[Any, \"T S\"\].",
         (
             "The current values for each jaxtyping PyTree structure annotation are as "
             "follows."
@@ -82,9 +82,9 @@ def test_dataclass_attribute(typecheck):
 
     matches = [
         "Type-check error whilst checking the parameters of M",
-        "The problem arose whilst typechecking argument 'z'.",
+        "The problem arose whilst typechecking parameter 'z'.",
         (
-            r"Called with arguments: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
+            r"Called with parameters: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
             r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
         ),
         (

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_pytree.py

@@ -26,7 +26,7 @@ import jax.random as jr
 import pytest
 
 import jaxtyping
-from jaxtyping import Array, Float, PyTree
+from jaxtyping import AnnotationError, Array, Float, PyTree
 
 from .helpers import make_mlp, ParamError
 
@@ -330,7 +330,7 @@ def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, get
 
     x1 = jr.normal(getkey(), (2,))
     y1 = jr.normal(getkey(), (2,))
-    with pytest.raises(ValueError, match="except when contained with structured"):
+    with pytest.raises(AnnotationError, match="except when contained with structured"):
         f(x1, y1)
 
 
@@ -340,5 +340,5 @@ def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, ge
         pass
 
     x1 = jr.normal(getkey(), (2,))
-    with pytest.raises(ValueError, match="ambiguous which PyTree"):
+    with pytest.raises(AnnotationError, match="ambiguous which PyTree"):
         f(x1)

test/test_serialisation.py

@@ -1,16 +1,25 @@
 import cloudpickle
 import numpy as np
+
+
+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)
+
+    if torch is not None:
+        w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
         cloudpickle.loads(w)

test/test_threading.py

@@ -39,8 +39,7 @@ class _ErrorableThread(threading.Thread):
 
 
 def test_threading_jaxtyped():
-    @jaxtyped
+    @jaxtyped(typechecker=typechecked)
-    @typechecked
     def add(x: Float[Array, "a b"], y: Float[Array, "a b"]) -> Float[Array, "a b"]:
         return x + y