Bump version

These flags are `JAXTYPING_DISABLE` and `JAXTYPING_REMOVE_TYPECHECKER_STACK`.

In addition, have now added warnings when using old-style double-decorator syntax, which also serves to guard against the easy mistake of
```python
@jaxtyped(typechecker)
def foo(...)
```
which actually decorates the `typechecker`, not `foo`.

.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
-    rev: 23.9.1
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.1.7
     hooks:
-    - id: black
-  - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.291'
-    hooks:
-      - id: ruff
-        args: ["--fix"]
+      - id: ruff  # linter
+        types_or: [ python, pyi, jupyter ]
+        args: [ --fix ]
+      - id: ruff-format  # formatter
+        types_or: [ python, pyi, jupyter ]

README.md

@@ -10,7 +10,7 @@ Type annotations **and runtime type-checking** for:
 ```python
 from jaxtyping import Array, Float, PyTree
 
-# Accepts floating-point 2D arrays with matching dimensions
+# Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:

docs/api/array.md

@@ -10,8 +10,9 @@ The shape should be a string of space-separated symbols, such as `"a b c d"`. Ea
 
 - `int`: fixed-size axis, e.g. `"28 28"`.
 - `str`: variable-size axis, e.g. `"channels"`.
-- A symbolic expression (without spaces!) in terms of other variable-size axes, e.g.  
-    `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`.
+- A symbolic expression in terms of other variable-size axes, e.g.  
+    `def remove_last(x: Float[Array, "dim"]) -> Float[Array, "dim-1"]`.  
+    Symbolic expressions must not use any spaces, otherwise each piece is treated as as a separate axis.
 
 When calling a function, variable-size axes and symbolic axes will be matched up across all arguments and checked for consistency. (See [Runtime type checking](./runtime-type-checking.md).)
 
@@ -19,11 +20,12 @@ When calling a function, variable-size axes and symbolic axes will be matched up
 
 In addition some modifiers can be applied:
 
-- Prepend `*` to a dimension to indicate that it can match multiple axes, e.g. `"*batch c h w"` will match zero or more batch axes.
-- Prepend `#` to a dimension to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g.  
+- Prepend `*` to an axis to indicate that it can match multiple axes, e.g. `"*batch"` will match zero or more batch axes.
+- Prepend `#` to an axis to indicate that it can be that size *or* equal to one -- i.e. broadcasting is acceptable, e.g.  
     `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
-- Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
+- 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/).)
 
 When using multiple modifiers, their order does not matter.
 
@@ -36,28 +38,38 @@ As a special case:
 - To denote a scalar shape use `""`, e.g. `Float[Array, ""]`.
 - To denote an arbitrary shape (and only check dtype) use `"..."`, e.g. `Float[Array, "..."]`.
 - You cannot have more than one use of multiple-axes, i.e. you can only use `...` or `*name` at most once in each array.
-- An example of broadcasting multiple dimensions:  
-    `def add(x: Float[Array, "*#foo"], y: Float[Array, "*#foo"]) -> Float[Array, "*#foo"]`.
 - A symbolic expression cannot be evaluated unless all of the axes sizes it refers to have already been processed. In practice this usually means that they should only be used in annotations for the return type, and only use axes declared in the arguments.
+- Symbolic expressions are evaluated in two stages: they are first evaluated as f-strings using the arguments of the function, and second are evaluated using the processed axis sizes. The f-string evaluation means that they can use local variables by enclosing them with curly braces, e.g. `{variable}`, e.g.
+    ```python
+    def full(size: int, fill: float) -> Float[Array, "{size}"]:
+        return jax.numpy.full((size,), fill)
+
+    class SomeClass:
+        some_value = 5
+
+        def full(self, fill: float) -> Float[Array, "{self.some_value}+3"]:
+            return jax.numpy.full((self.some_value + 3,), fill)
+    ```
 
 ## Dtype
 
 The dtype should be any one of (all imported from `jaxtyping`):
 
 - Any dtype at all: `Shaped`
-  - Boolean: `Bool`
-  - PRNG key: `Key`
-  - Any integer, unsigned integer, floating, or complex: `Num`
-    - Any floating or complex: `Inexact`
-      - Any floating point: `Float`
-        - Of particular precision: `BFloat16`, `Float16`, `Float32`, `Float64`
-      - Any complex: `Complex`
-        - Of particular precision: `Complex64`, `Complex128`
-    - Any integer or unsigned intger: `Integer`
-      - Any unsigned integer: `UInt`
-        - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
-      - Any signed integer: `Int`
-        - Of particular precision: `Int8`, `Int16`, `Int32`, `Int64`
+    - Boolean: `Bool`
+    - PRNG key: `Key`
+    - Any integer, unsigned integer, floating, or complex: `Num`
+        - Any floating or complex: `Inexact`
+            - Any floating point: `Float`
+                - Of particular precision: `BFloat16`, `Float16`, `Float32`, `Float64`
+            - Any complex: `Complex`
+                - Of particular precision: `Complex64`, `Complex128`
+        - Any integer or unsigned intger: `Integer`
+            - Any unsigned integer: `UInt`
+                - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
+            - Any signed integer: `Int`
+                - Of particular precision: `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
 ```python

docs/api/pytree.md

@@ -11,4 +11,36 @@
 
 ---
 
+## Path-dependent shapes
+
+The prefix `?` may be used to indicate that the axis size can depend on which leaf of a PyTree the array is at. For example:
+```python
+def f(
+    x: PyTree[Shaped[Array, "?foo"], "T"],
+    y: PyTree[Shaped[Array, "?foo"], "T"],
+):
+    pass
+```
+The above demands that `x` and `y` have matching PyTree structures (due to the `T` annotation), and that their leaves must all be one-dimensional arrays, *and that the corresponding pairs of leaves in `x` and `y` must have the same size as each other*.
+
+Thus the following is allowed:
+```python
+x0 = jnp.arange(3)
+x1 = jnp.arange(5)
+
+y0 = jnp.arange(3) + 1
+y1 = jnp.arange(5) + 1
+
+f((x0, x1), (y0, y1))  # x0 matches y0, and x1 matches y1. All good!
+```
+
+But this is not:
+```python
+f((x1, x1), (y0, y1))  # x1 does not have a size matching y0!
+```
+
+Internally, all that is happening is that `foo` is replaced with `0foo` for the first leaf, `1foo` for the next leaf, etc., so that each leaf gets a unique version of the name.
+
+---
+
 Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.

docs/api/runtime-type-checking.md

@@ -2,27 +2,40 @@
 
 (See the [FAQ](../faq.md) for details on static type checking.)
 
-Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed at trace-time only, and will not impact runtime performance.
+Runtime type checking **synergises beautifully with `jax.jit`!** All shape checks will be performed only whilst tracing, and will not impact runtime performance.
 
-Runtime type-checking should be performed using a library like [typeguard](https://github.com/agronholm/typeguard) or [beartype](https://github.com/beartype/beartype).
+There are two approaches: either use [`jaxtyping.jaxtyped`][] to typecheck a single function, or [`jaxtyping.install_import_hook`][] to typecheck a whole codebase.
 
-The types provided by `jaxtyping`, e.g. `Float[Array, "batch channels"]`, are all compatible with `isinstance` checks, e.g. `isinstance(x, Float[Array, "batch channels"])`. This means that jaxtyping should be compatible with all runtime type checkers out-of-the-box.
+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.)
 
-Some additional context is needed to ensure consistency between multiple argments (i.e. that shapes match up between arrays). For this, you can use either `jaxtyping.jaxtyped` to add this capability to a single function, or `jaxtyping.install_import_hook` to add this capability to a whole codebase. If either are too much magic for you, you can safely use neither and have just single-argument type checking.
+---
 
 ::: jaxtyping.jaxtyped
 
 ---
 
-It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
-(Not to mention that double-decorators everywhere are a bit ugly.)
-
-The easier option is usually to use the import hook.
-
 ::: jaxtyping.install_import_hook
 
 ---
 
+#### Pytest hook
+
+The import hook can be installed at test-time only, as a pytest hook. From the command line the syntax is:
+```
+pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
+```
+or in `pyproject.toml`:
+```toml
+[tool.pytest.ini_options]
+addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
+```
+or in `pytest.ini`:
+```ini
+[pytest]
+addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
+```
+This example will apply the import hook to all modules whose names start with either `foo` or `bar.baz`. The typechecker used in this example is `beartype.beartype`.
+
 #### IPython extension
 
 If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
@@ -32,3 +45,7 @@ import jaxtyping
 %jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
 ```
 Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.
+
+#### Other runtime type-checking libraries
+
+Beartype and typeguard happen to be the two most popular runtime type-checking libraries (at least at time of writing), but jaxtyping should be compatible with all runtime type checkers out-of-the-box. The runtime type-checking library just needs to provide a type-checking decorator (analgous to `beartype.beartype` or `typeguard.typechecked`), and perform `isinstance` checks against jaxtyping's types.

docs/index.md

@@ -24,7 +24,7 @@ The annotations provided by jaxtyping are compatible with runtime type-checking
 ```python
 from jaxtyping import Array, Float, PyTree
 
-# Accepts floating-point 2D arrays with matching dimensions
+# Accepts floating-point 2D arrays with matching axes
 def matrix_multiply(x: Float[Array, "dim1 dim2"],
                     y: Float[Array, "dim2 dim3"]
                   ) -> Float[Array, "dim1 dim3"]:

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

@@ -29,7 +29,12 @@ from ._array_types import (
     has_jax,
     set_array_name_format as set_array_name_format,
 )
+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
 
@@ -65,6 +70,8 @@ elif has_jax:
 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 +91,7 @@ if typing.TYPE_CHECKING:
         Integer as Integer,
         Key as Key,
         Num as Num,
+        Real as Real,
         Shaped as Shaped,
         UInt as UInt,
         UInt8 as UInt8,
@@ -110,6 +118,7 @@ else:
         Int64 as Int64,
         Integer as Integer,
         Num as Num,
+        Real as Real,
         Shaped as Shaped,
         UInt as UInt,
         UInt8 as UInt8,
@@ -119,8 +128,13 @@ else:
     )
 
     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:
@@ -178,9 +192,11 @@ elif has_jax:
 
 # Conveniences
 if typing.TYPE_CHECKING:
-    from jax.random import PRNGKeyArray as PRNGKeyArray
-
-    from ._indirection import Scalar as Scalar, ScalarLike as ScalarLike
+    from ._indirection import (
+        PRNGKeyArray as PRNGKeyArray,
+        Scalar as Scalar,
+        ScalarLike as ScalarLike,
+    )
 elif has_jax:
     from ._array_types import Scalar, ScalarLike  # noqa: F401
 

jaxtyping/_array_types.py

@@ -27,7 +27,13 @@ from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
-from ._decorator import storage
+from ._errors import AnnotationError
+from ._storage import (
+    get_shape_memo,
+    get_treeflatten_memo,
+    get_treepath_memo,
+    set_shape_memo,
+)
 
 
 try:
@@ -68,15 +74,17 @@ class _DimType(enum.Enum):
 
 
 class _NamedDim:
-    def __init__(self, name, broadcastable):
+    def __init__(self, name, broadcastable, treepath):
         self.name = name
         self.broadcastable = broadcastable
+        self.treepath = treepath
 
 
 class _NamedVariadicDim:
-    def __init__(self, name, broadcastable):
+    def __init__(self, name, broadcastable, treepath):
         self.name = name
         self.broadcastable = broadcastable
+        self.treepath = treepath
 
 
 class _FixedDim:
@@ -86,8 +94,8 @@ class _FixedDim:
 
 
 class _SymbolicDim:
-    def __init__(self, expr, broadcastable):
-        self.expr = expr
+    def __init__(self, elem, broadcastable):
+        self.elem = elem
         self.broadcastable = broadcastable
 
 
@@ -104,8 +112,9 @@ _AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim
 
 def _check_dims(
     cls_dims: list[_AbstractDim],
-    obj_shape: tuple[int],
+    obj_shape: tuple[int, ...],
     single_memo: dict[str, int],
+    arg_memo: dict[str, Any],
 ) -> bool:
     assert len(cls_dims) == len(obj_shape)
     for cls_dim, obj_size in zip(cls_dims, obj_shape):
@@ -118,55 +127,45 @@ def _check_dims(
                 return False
         elif type(cls_dim) is _SymbolicDim:
             try:
+                # Support f-string syntax.
+                # https://stackoverflow.com/a/53671539/22545467
+                elem = eval(f"f'{cls_dim.elem}'", arg_memo.copy())
                 # Make a copy to avoid `__builtins__` getting added as a key.
-                eval_size = eval(cls_dim.expr, single_memo.copy())
+                eval_size = eval(elem, single_memo.copy())
             except NameError as e:
-                raise NameError(
-                    f"Cannot process symbolic dimension '{cls_dim.expr}' as some "
-                    "dimension names have not been processed. In practice you should "
-                    "usually only use symbolic dimensions in annotations for return "
-                    "types, referring only to dimensions annotated for arguments."
+                raise AnnotationError(
+                    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
             if eval_size != obj_size:
                 return False
         else:
             assert type(cls_dim) is _NamedDim
+            if cls_dim.treepath:
+                name = get_treepath_memo() + cls_dim.name
+            else:
+                name = cls_dim.name
             try:
-                cls_size = single_memo[cls_dim.name]
+                cls_size = single_memo[name]
             except KeyError:
-                single_memo[cls_dim.name] = obj_size
+                single_memo[name] = obj_size
             else:
                 if cls_size != obj_size:
                     return False
     return True
 
 
-def _is_jax_extended_dtype(dtype: Any) -> bool:
-    if not has_jax:
-        return False
-    try:
-        is_dtype = issubclass(dtype, jax.numpy.generic)
-    except TypeError:
-        # `dtype` not a class
-        return False
-    else:
-        if is_dtype:
-            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
-                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
-            else:  # jax<=0.4.13
-                return jax.core.is_opaque_dtype(dtype)
-        else:
-            return False
-
-
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
+        if get_treeflatten_memo():
+            return True
 
-        if _is_jax_extended_dtype(obj.dtype):
-            dtype = str(obj.dtype)
-        elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
+        if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
             dtype = obj.dtype.type.__name__
         elif hasattr(obj.dtype, "as_numpy_dtype"):
@@ -178,7 +177,7 @@ class _MetaAbstractArray(type):
             if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
                 dtype = repr_dtype[1]
             else:
-                raise RuntimeError(
+                raise AnnotationError(
                     "Unrecognised array/tensor type to extract dtype from"
                 )
 
@@ -196,45 +195,37 @@ class _MetaAbstractArray(type):
             if not in_dtypes:
                 return False
 
-        no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
-
-        if no_temp_memo:
-            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
-            # Make a copy so we don't mutate the original memo during the shape check.
-            single_memo = single_memo.copy()
-            variadic_memo = variadic_memo.copy()
-            variadic_broadcast_memo = variadic_broadcast_memo.copy()
-        else:
-            # `isinstance` happening outside any @jaxtyped decorators, e.g. at the
-            # global scope. In this case just create a temporary memo, since we're not
-            # going to be comparing against any stored values anyway.
-            single_memo = {}
-            variadic_memo = {}
-            variadic_broadcast_memo = {}
-
-        if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo):
-            # We update the memo every time we successfully pass a shape check
-            if no_temp_memo:
-                storage.memo_stack[-1] = (
-                    single_memo,
-                    variadic_memo,
-                    variadic_broadcast_memo,
-                )
+        single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        arg_memo_bak = arg_memo.copy()
+        try:
+            check = cls._check_shape(obj, single_memo, variadic_memo, arg_memo)
+        except Exception:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            raise
+        if check:
             return True
         else:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
             return False
 
     def _check_shape(
         cls,
         obj,
         single_memo: dict[str, int],
-        variadic_memo: dict[str, tuple[int, ...]],
-        variadic_broadcast_memo: dict[str, list[tuple[int, ...]]],
+        variadic_memo: dict[str, tuple[bool, tuple[int, ...]]],
+        arg_memo: dict[str, Any],
     ):
         if cls.index_variadic is None:
             if obj.ndim != len(cls.dims):
                 return False
-            return _check_dims(cls.dims, obj.shape, single_memo)
+            return _check_dims(cls.dims, obj.shape, single_memo, arg_memo)
         else:
             if obj.ndim < len(cls.dims) - 1:
                 return False
@@ -242,10 +233,10 @@ class _MetaAbstractArray(type):
             j = -(len(cls.dims) - i - 1)
             if j == 0:
                 j = None
-            if not _check_dims(cls.dims[:i], obj.shape[:i], single_memo):
+            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
+                cls.dims[j:], obj.shape[j:], single_memo, arg_memo
             ):
                 return False
             variadic_dim = cls.dims[i]
@@ -253,30 +244,40 @@ class _MetaAbstractArray(type):
                 return True
             else:
                 assert type(variadic_dim) is _NamedVariadicDim
-                variadic_name = variadic_dim.name
+                if variadic_dim.treepath:
+                    name = get_treepath_memo() + variadic_dim.name
+                else:
+                    name = variadic_dim.name
+                broadcastable = variadic_dim.broadcastable
                 try:
-                    if variadic_dim.broadcastable:
-                        variadic_shapes = variadic_broadcast_memo[variadic_name]
-                    else:
-                        variadic_shape = variadic_memo[variadic_name]
+                    prev_broadcastable, prev_shape = variadic_memo[name]
                 except KeyError:
-                    if variadic_dim.broadcastable:
-                        variadic_broadcast_memo[variadic_name] = [obj.shape[i:j]]
-                    else:
-                        variadic_memo[variadic_name] = obj.shape[i:j]
+                    variadic_memo[name] = (broadcastable, obj.shape[i:j])
                     return True
                 else:
-                    if variadic_dim.broadcastable:
-                        new_shape = obj.shape[i:j]
-                        for existing_shape in variadic_shapes:
-                            try:
-                                np.broadcast_shapes(new_shape, existing_shape)
-                            except ValueError:
-                                return False
-                        variadic_shapes.append(new_shape)
-                        return True
+                    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
+                        if not broadcastable and broadcast_shape != new_shape:
+                            return False
+                        variadic_memo[name] = (broadcastable, broadcast_shape)
                     else:
-                        return variadic_shape == obj.shape[i:j]
+                        if broadcastable:
+                            try:
+                                broadcast_shape = np.broadcast_shapes(
+                                    new_shape, prev_shape
+                                )
+                            except ValueError:  # not broadcastable e.g. (3, 4) and (5,)
+                                return False
+                            if broadcast_shape != prev_shape:
+                                return False
+                        else:
+                            if new_shape != prev_shape:
+                                return False
+                    return True
             assert False
 
 
@@ -289,8 +290,11 @@ def _make_metaclass(base_metaclass):
 
 
 def _check_scalar(dtype, dtypes, dims):
-    if len(dims) != 0:
-        return dims == (_anonymous_variadic_dim,)
+    for dim in dims:
+        if dim is not _anonymous_variadic_dim and not isinstance(
+            dim, _NamedVariadicDim
+        ):
+            return False
     return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
 
 
@@ -330,28 +334,30 @@ def _make_array(array_type, dim_str, dtypes, name):
         if "," in elem and "(" not in elem:
             # Common mistake.
             # Disable in the case that there's brackets to allow for function calls,
-            # e.g. `min(foo,bar)`, in symbolic dimensions.
-            raise ValueError("Dimensions should be separated with spaces, not commas")
+            # e.g. `min(foo,bar)`, in symbolic axes.
+            raise ValueError("Axes should be separated with spaces, not commas")
         if elem.endswith("#"):
             raise ValueError(
-                "As of jaxtyping v0.1.0, broadcastable dimensions are now denoted "
+                "As of jaxtyping v0.1.0, broadcastable axes are now denoted "
                 "with a # at the start, rather than at the end"
             )
 
         if "..." in elem:
             if elem != "...":
                 raise ValueError(
-                    "Anonymous multiple dimension '...' must be used on its own; "
+                    "Anonymous multiple axes '...' must be used on its own; "
                     f"got {elem}"
                 )
             broadcastable = False
             variadic = True
             anonymous = True
+            treepath = False
             dim_type = _DimType.named
         else:
             broadcastable = False
             variadic = False
             anonymous = False
+            treepath = False
             while True:
                 if len(elem) == 0:
                     # This branch needed as just `_` is valid
@@ -369,7 +375,7 @@ def _make_array(array_type, dim_str, dtypes, name):
                     if variadic:
                         raise ValueError(
                             "Do not use * twice to denote accepting multiple "
-                            "dimensions, e.g. `**foo` is not allowed"
+                            "axes, e.g. `**foo` is not allowed"
                         )
                     variadic = True
                     elem = elem[1:]
@@ -381,6 +387,14 @@ def _make_array(array_type, dim_str, dtypes, name):
                         )
                     anonymous = True
                     elem = elem[1:]
+                elif first_char == "?":
+                    if treepath:
+                        raise ValueError(
+                            "Do not use ? twice to denote dependence on location "
+                            "within a PyTree, e.g. `??foo` is not allowed"
+                        )
+                    treepath = True
+                    elem = elem[1:]
                 # Allow e.g. `foo=4` as an alternate syntax for just `4`, so that one
                 # can write e.g. `Float[Array, "rows=3 cols=4"]`
                 elif elem.count("=") == 1:
@@ -400,28 +414,33 @@ def _make_array(array_type, dim_str, dtypes, name):
         if variadic:
             if index_variadic is not None:
                 raise ValueError(
-                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
-                    "more than once"
+                    "Cannot use variadic specifiers (`*name` or `...`) "
+                    "more than once."
                 )
             index_variadic = index
 
         if dim_type is _DimType.fixed:
             if variadic:
                 raise ValueError(
-                    "Cannot have a fixed axis bind to multiple dimensions, e.g. "
-                    "`*4` is not allowed"
+                    "Cannot have a fixed axis bind to multiple axes, e.g. "
+                    "`*4` is not allowed."
                 )
             if anonymous:
                 raise ValueError(
-                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not " "allowed"
+                    "Cannot have a fixed axis be anonymous, e.g. `_4` is not allowed."
+                )
+            if treepath:
+                raise ValueError(
+                    "Cannot have a fixed axis have tree-path dependence, e.g. `?4` is "
+                    "not allowed."
                 )
             elem = _FixedDim(elem, broadcastable)
         elif dim_type is _DimType.named:
             if anonymous:
                 if broadcastable:
                     raise ValueError(
-                        "Cannot have a dimension be both anonymous and "
-                        "broadcastable, e.g. `#_` is not allowed"
+                        "Cannot have an axis be both anonymous and "
+                        "broadcastable, e.g. `#_` is not allowed."
                     )
                 if variadic:
                     elem = _anonymous_variadic_dim
@@ -429,22 +448,26 @@ def _make_array(array_type, dim_str, dtypes, name):
                     elem = _anonymous_dim
             else:
                 if variadic:
-                    elem = _NamedVariadicDim(elem, broadcastable)
+                    elem = _NamedVariadicDim(elem, broadcastable, treepath)
                 else:
-                    elem = _NamedDim(elem, broadcastable)
+                    elem = _NamedDim(elem, broadcastable, treepath)
         else:
             assert dim_type is _DimType.symbolic
             if anonymous:
                 raise ValueError(
-                    "Cannot have a symbolic dimension be anonymous, e.g. "
+                    "Cannot have a symbolic axis be anonymous, e.g. "
                     "`_foo+bar` is not allowed"
                 )
             if variadic:
                 raise ValueError(
-                    "Cannot have symbolic multiple-dimensions, e.g. "
+                    "Cannot have symbolic multiple-axes, e.g. "
                     "`*foo+bar` is not allowed"
                 )
-            elem = compile(elem, "<string>", "eval")
+            if treepath:
+                raise ValueError(
+                    "Cannot have a symbolic axis with tree-path dependence, e.g. "
+                    "`?foo+bar` is not allowed"
+                )
             elem = _SymbolicDim(elem, broadcastable)
         dims.append(elem)
     dims = tuple(dims)
@@ -472,6 +495,16 @@ def _make_array(array_type, dim_str, dtypes, name):
             return array_type
         else:
             return _not_made
+    elif array_type is np.bool_:
+        if _check_scalar("bool", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is np.generic or array_type is np.number:
+        if _check_scalar("", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
     if issubclass(array_type, AbstractArray):
         if dtypes is _any_dtype:
             dtypes = array_type.dtypes
@@ -488,7 +521,7 @@ def _make_array(array_type, dim_str, dtypes, name):
                 index_variadic = array_type.index_variadic + len(dims)
             else:
                 raise ValueError(
-                    "Cannot use multiple-dimension specifiers (`*name` or `...`) "
+                    "Cannot use variadic specifiers (`*name` or `...`) "
                     "in both the original array and the extended array"
                 )
         dims = dims + array_type.dims
@@ -525,7 +558,7 @@ def _make_array(array_type, dim_str, dtypes, name):
 
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
-        raise RuntimeError(
+        raise AnnotationError(
             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, "..."]`.'
@@ -539,6 +572,7 @@ class _MetaAbstractDtype(type):
                 "Ellipsis can be used to accept any shape: `Float[Array, '...']`."
             )
         array_type, dim_str = item
+        dim_str = dim_str.strip()
         del item
         if typing.get_origin(array_type) in _union_types:
             out = [
@@ -604,6 +638,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
         cls.dtypes = dtypes
 
 
+_prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
 _uint8 = "uint8"
@@ -666,13 +701,13 @@ Integer = _make_dtype(uints + ints, "Integer")
 Float = _make_dtype(floats, "Float")
 Complex = _make_dtype(complexes, "Complex")
 Inexact = _make_dtype(floats + complexes, "Inexact")
+Real = _make_dtype(floats + uints + ints, "Real")
 Num = _make_dtype(uints + ints + floats + complexes, "Num")
 
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 if has_jax:
-    _key_regex = re.compile(r"^key<\w+>$")
-    Key = _make_dtype(_key_regex, "Key")
+    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"]]

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

@@ -20,10 +20,10 @@
 import dataclasses
 import functools as ft
 import inspect
-import threading
-import types
-import weakref
-from typing import get_args, get_origin
+import itertools as it
+import sys
+import warnings
+from typing import Any, get_args, get_origin, get_type_hints, overload
 
 
 try:
@@ -34,125 +34,458 @@ else:
     traceback_util.register_exclusion(__file__)
 
 
-storage = threading.local()
+from ._config import config
+from ._errors import AnnotationError, TypeCheckError
+from ._storage import pop_shape_memo, push_shape_memo
 
 
-_jaxtyped_fns = weakref.WeakSet()
+class _Sentinel:
+    def __repr__(self):
+        return "sentinel"
 
 
-def jaxtyped(fn):
-    """Used in conjunction with a runtime type checker. Decorate a function with this to
-    have shapes checked for consistency across multiple arguments.
+_sentinel = _Sentinel()
 
-    Note that `@jaxtyped` is applied above the type checker.
+
+@overload
+def jaxtyped(*, typechecker=_sentinel):
+    ...
+
+
+@overload
+def jaxtyped(fn, *, typechecker=_sentinel):
+    ...
+
+
+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.
 
     !!! Example
 
         ```python
         # Import both the annotation and the `jaxtyped` decorator from `jaxtyping`
-        from jaxtyping import Array, Float32, jaxtyped
+        from jaxtyping import Array, Float, jaxtyped
 
         # Use your favourite typechecker: usually one of the two lines below.
         from typeguard import typechecked as typechecker
         from beartype import beartype as typechecker
 
-        # Write your function. @jaxtyped must be applied above @typechecker!
-        @jaxtyped
-        @typechecker
-        def batch_outer_product(x: Float32[Array, "b c1"],
-                                y: Float32[Array, "b c2"]
-                              ) -> Float32[Array, "b c1 c2"]:
+        # Type-check a function
+        @jaxtyped(typechecker=typechecker)
+        def batch_outer_product(x: Float[Array, "b c1"],
+                                y: Float[Array, "b c2"]
+                              ) -> Float[Array, "b c1 c2"]:
             return x[:, :, None] * y[:, None, :]
+
+        # Type-check a dataclass
+        @jaxtyped(typechecker=typechecker)
+        @dataclass
+        class MyDataclass:
+            x: int
+            y: Float[Array "b c"]
         ```
 
-    **Notes for advanced users**
+    **Arguments:**
 
-    Put precisely, all `isinstance` shape checks are scoped to the thread-local dynamic
-    context of a `jaxtyped` call. A new dynamic context will allow different dimensions
-    sizes to be bound to the same name. After this new dynamic context is finished
-    then the old one is returned to.
+    - `fn`: The function or dataclass to decorate.
+    - `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):
+            pass
 
-    For example, this means you could leave off the `@jaxtyped` decorator to enforce
-    that this function use the same axis sizes as the function it was called from.
+        f("a string is not an integer")  # this line should raise an exception
+        ```
+        Common choices are `typechecker=beartype.beartype` or
+        `typechecker=typeguard.typechecked`. Can also be set as `typechecker=None` to
+        skip automatic runtime type-checking, but still support manual `isinstance`
+        checks inside the function body:
+        ```python
+        @jaxtyped(typechecker=None)
+        def f(x):
+            assert isinstance(x, Float[Array, "batch channel"])
+        ```
 
-    Likewise, this means you can use `isinstance` checks inside a function body
-    and have them contribute to the same collection of consistency checks performed
-    by a typechecker against its arguments. (Or even forgo a typechecker that analyses
-    arguments, and instead just do your own manual `isinstance` checks.)
+    **Returns:**
 
-    Only `isinstance` checks that pass will contribute to the store of axis name-size
-    pairs; those that fail will not. As such it is safe to write e.g.
-    `assert not isinstance(x, Float32[Array, "foo"])`.
+    If `fn` is a function (including a `staticmethod`, `classmethod`, or `property`),
+    then a wrapped function is returned.
+
+    If `fn` is a dataclass, then `fn` is returned directly, and additionally its
+    `__init__` method is wrapped and modified in-place.
+
+    !!! Info "Old syntax"
+
+        jaxtyping previously (before v0.2.24) recommended using this double-decorator
+        syntax:
+        ```python
+        @jaxtyped
+        @typechecker
+        def f(...): ...
+        ```
+        This is still supported, but will now raise a warning recommending the
+        `jaxtyped(typechecker=typechecker)` syntax discussed above. (Which will produce
+        easier-to-debug error messages: under the hood, the new syntax more carefully
+        manipulates the typechecker so as to determine where a type-check error arises.)
+
+    ??? Info "Notes for advanced users"
+
+        **Dynamic contexts:**
+
+        Put precisely, the axis names in e.g. `Float[Array, "batch channels"]` and the
+        structure names in e.g. `PyTree[int, "T"]` are all scoped to the thread-local
+        dynamic context of a `jaxtyped`-wrapped function. If from within that function
+        we then call another `jaxtyped`-wrapped function, then a new context is pushed
+        to the stack. The axis sizes and PyTree structures of this inner function will
+        then not be compared against the axis sizes and PyTree structures of the outer
+        function. After the inner function returns then this inner context is popped
+        from the stack, and the previous context is returned to.
+
+        **isinstance:**
+
+        Binding of a value against a name is done with an `isinstance` check, for
+        example `isinstance(jnp.zeros((3, 4)), Float[Array, "dim1 dim2"])` will bind
+        `dim1=3` and `dim2=4`. In practice these `isinstance` checks are usually done by
+        the run-time typechecker `typechecker` that is supplied as an argument.
+
+        This can also be done manually: add `isinstance` checks inside a function body
+        and they will contribute to the same collection of consistency checks as are
+        performed by the typechecker on the arguments and return values. (Or you can
+        forgo such a typechecker altogether -- i.e. `typechecker=None` -- and only do
+        your own manual `isinstance` checks.)
+
+        Only `isinstance` checks that pass will contribute to the store of values; those
+        that fail will not. As such it is safe to write e.g.
+        `assert not isinstance(x, Float32[Array, "foo"])`.
+
+        **Decoupling contexts from function calls:**
+
+        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.
+
+        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"])
+        ```
+        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.
     """
-    if type(fn) is types.FunctionType and fn in _jaxtyped_fns:
-        return fn
-    elif inspect.isclass(fn):  # allow decorators on class definitions
-        if dataclasses.is_dataclass(fn):
-            init = jaxtyped(fn.__init__)
-            fn.__init__ = init
-            return fn
-        else:
+
+    # First handle the `jaxtyped("context")` usage, which is a special case.
+    if fn == "context":
+        if typechecker is not _sentinel:
             raise ValueError(
-                "jaxtyped may only be added as a class decorator to dataclasses"
+                "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 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
+            # type annotations. There may be a custom user `__init__`, or a
+            # dataclass-generated `__init__` used alongside
+            # `equinox.field(converter=...)`
+
+            init = fn.__init__
+
+            @ft.wraps(init)
+            def __init__(self, *args, **kwargs):
+                init(self, *args, **kwargs)
+                # `fn.__init__` is late-binding to the `__init__` function that
+                # we're in now. (Or to someone else's monkey-patch.) Either way,
+                # this checks that we're in the "top-level" `__init__`, and not one
+                # that is being called via `super()`. We don't want to trigger too
+                # early, before all fields have been assigned.
+                #
+                # We're not checking `if self.__class__ is fn` because Equinox
+                # replaces the with a defrozen version of itself during `__init__`,
+                # so the check wouldn't trigger.
+                #
+                # We're not doing this check by adding it to the end of the
+                # metaclass `__call__`, because Python doesn't allow you
+                # monkey-patch metaclasses.
+                if self.__class__.__init__ is fn.__init__:
+                    _check_dataclass_annotations(self, typechecker)
+
+            fn.__init__ = __init__
+        return fn
     # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin
     # ones. Unfortunately that means returning a class instance with a __get__ method,
     # and that turns out to break loads of other things. See beartype issue #211 and
     # jaxtyping issue #71.
     elif isinstance(fn, classmethod):
-        return classmethod(jaxtyped(fn.__func__))
+        return classmethod(jaxtyped(fn.__func__, typechecker=typechecker))
     elif isinstance(fn, staticmethod):
-        return staticmethod(jaxtyped(fn.__func__))
+        return staticmethod(jaxtyped(fn.__func__, typechecker=typechecker))
     elif isinstance(fn, property):
         if fn.fget is None:
             fget = None
         else:
-            fget = jaxtyped(fn.fget)
+            fget = jaxtyped(fn.fget, typechecker=typechecker)
         if fn.fset is None:
             fset = None
         else:
-            fset = jaxtyped(fn.fset)
+            fset = jaxtyped(fn.fset, typechecker=typechecker)
         if fn.fdel is None:
             fdel = None
         else:
-            fdel = jaxtyped(fn.fdel)
+            fdel = jaxtyped(fn.fdel, typechecker=typechecker)
         return property(fget=fget, fset=fset, fdel=fdel)
     else:
+        if typechecker is None:
+            # Probably being used in the old style as
+            # ```
+            # @jaxtyped
+            # @typechecker
+            # def foo(x: int): ...
+            # ```
+            # in which case make a best-effort attempt to add shape information for any
+            # type errors.
 
-        @ft.wraps(fn)
-        def wrapped_fn(*args, **kwargs):
+            signature = inspect.signature(fn)
+
+            @ft.wraps(fn)
+            def wrapped_fn(*args, **kwargs):  # pyright: ignore
+                bound = signature.bind(*args, **kwargs)
+                memos = push_shape_memo(bound.arguments)
+                try:
+                    return fn(*args, **kwargs)
+                except Exception as e:
+                    if sys.version_info >= (3, 11) and _no_jaxtyping_note(e):
+                        shape_info = _exc_shape_info(memos)
+                        if shape_info != "":
+                            msg = (
+                                "The preceding error occurred within the scope of a "
+                                "`jaxtyping.jaxtyped` function, and may be due to a "
+                                "typecheck error. "
+                            )
+                            e.add_note(_jaxtyping_note_str(_spacer + msg + shape_info))
+                    raise
+                finally:
+                    pop_shape_memo()
+
+        else:
+            # New-style
+            # ```
+            # @jaxtyped(typechecker=typechecker)
+            # def foo(x: int): ...
+            # ```
+            # in which case we can do a better job reporting errors.
+
+            full_signature = inspect.signature(fn)
             try:
-                memo_stack = storage.memo_stack
-            except AttributeError:
-                memo_stack = storage.memo_stack = []
-            memo_stack.append(({}, {}, {}))
-            try:
-                return fn(*args, **kwargs)
-            finally:
-                memo_stack.pop()
+                destring_annotations = get_type_hints(fn, include_extras=True)
+            except NameError:
+                # Best-effort attempt to destringify annotations.
+                pass
+            else:
+                new_params = []
+                for p_name, p_value in full_signature.parameters.items():
+                    p_annotation = destring_annotations.get(p_name, p_value.annotation)
+                    p_value = p_value.replace(annotation=p_annotation)
+                    new_params.append(p_value)
+                return_annotation = destring_annotations.get(
+                    "return", full_signature.return_annotation
+                )
+                full_signature = full_signature.replace(
+                    parameters=new_params, return_annotation=return_annotation
+                )
+
+            param_signature = full_signature.replace(
+                return_annotation=inspect.Signature.empty
+            )
+            module = getattr(fn, "__module__", "generated")
+
+            full_fn, output_name = _make_fn_with_signature(
+                "check_return", full_signature, module, output=True
+            )
+            full_fn = typechecker(full_fn)
+
+            param_fn = _make_fn_with_signature(
+                "check_params", param_signature, module, output=False
+            )
+            param_fn = typechecker(param_fn)
+
+            @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)
+
+                memos = push_shape_memo(bound.arguments)
+                try:
+                    # First type-check just the parameters before the function is
+                    # called.
+                    try:
+                        param_fn(*args, **kwargs)
+                    except AnnotationError:
+                        raise
+                    except Exception as e:
+                        argmsg = _get_problem_arg(
+                            param_signature,
+                            args,
+                            kwargs,
+                            bound.arguments,
+                            module,
+                            typechecker,
+                        )
+                        try:
+                            name = fn.__name__
+                        except AttributeError:
+                            name = fn.__class__.__name__
+                        param_values = _pformat(bound.arguments, short_self=True)
+                        param_hints = _remove_typing(param_signature)
+                        msg = (
+                            "Type-check error whilst checking the parameters of "
+                            f"{name}.{argmsg}\n"
+                            "----------------------\n"
+                            f"Called with parameters: {param_values}\n"
+                            f"Parameter annotations: {param_hints}.\n"
+                            + _exc_shape_info(memos)
+                        )
+                        if config.jaxtyping_remove_typechecker_stack:
+                            raise TypeCheckError(msg) from None
+                        else:
+                            raise TypeCheckError(msg) from e
+
+                    # Actually call the function.
+                    out = fn(*args, **kwargs)
+
+                    if full_signature.return_annotation is not inspect.Signature.empty:
+                        # Now type-check the return value. We need to include the
+                        # parameters in the type-checking here in case there are any
+                        # type variables shared across the parameters and return.
+                        #
+                        # Incidentally this does mean that if `fn` mutates its arguments
+                        # so that they no longer satisfy their type annotations, this
+                        # will throw an error here. But that's like, super weird, so
+                        # don't do that. An error in that scenario is probably still
+                        # desirable.
+                        #
+                        # There is a small performance concern here when used in
+                        # non-jit'd contexts, like PyTorch, due to the duplicate
+                        # checking of the parameters. Unfortunately there doesn't seem
+                        # to be a way around that, so c'est la vie.
+                        kwargs[output_name] = out
+                        try:
+                            full_fn(*args, **kwargs)
+                        except AnnotationError:
+                            raise
+                        except Exception as e:
+                            try:
+                                name = fn.__name__
+                            except AttributeError:
+                                name = fn.__class__.__name__
+                            param_values = _pformat(bound.arguments, short_self=True)
+                            return_value = _pformat(out, short_self=False)
+                            param_hints = _remove_typing(param_signature)
+                            return_hint = _remove_typing(
+                                full_signature.return_annotation
+                            )
+                            if return_hint.startswith(
+                                "<class '"
+                            ) and return_hint.endswith("'>"):
+                                return_hint = return_hint[8:-2]
+                            msg = (
+                                "Type-check error whilst checking the return value "
+                                f"of {name}.\n"
+                                f"Actual value: {return_value}\n"
+                                f"Expected type: {return_hint}.\n"
+                                "----------------------\n"
+                                f"Called with parameters: {param_values}\n"
+                                f"Parameter annotations: {param_hints}.\n"
+                                + _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
 
 
-@jaxtyped
+class _JaxtypingContext:
+    def __enter__(self):
+        push_shape_memo({})
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        pop_shape_memo()
+
+
 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.
+    `beartype.beartype` or `typeguard.typechecked`.
+    """
+    parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)]
+    values = {}
     for field in dataclasses.fields(self):
-        for kls in self.__class__.__mro__:
-            try:
-                annotation = kls.__annotations__[field.name]
-            except KeyError:
-                pass
-            else:
-                break
-        else:
-            raise TypeError
+        annotation = field.type
         if isinstance(annotation, str):
-            # Don't support stringified annotations. These are basically impossible to
+            # Don't check stringified annotations. These are basically impossible to
             # resolve correctly, so just skip them.
-            # This does mean that annotations like `type["Foo"]` will just fail. There
-            # doesn't seem to be any way to even detect a partially-stringified
-            # annotation.
             continue
         if get_origin(annotation) is type:
             args = get_args(annotation)
@@ -163,66 +496,302 @@ def _check_dataclass_annotations(self, typechecker):
                 # https://github.com/patrick-kidger/equinox/pull/543
                 continue
         try:
-            value = getattr(self, field.name)
+            value = getattr(self, field.name)  # noqa: F841
         except AttributeError:
             continue  # allow uninitialised fields, which are allowed on dataclasses
 
-        @typechecker
-        def typecheck(x: annotation):
-            pass
+        parameters.append(
+            inspect.Parameter(
+                field.name,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=field.type,
+            )
+        )
+        values[field.name] = value
 
-        typecheck(value)
+    signature = inspect.Signature(parameters)
+    module = self.__class__.__module__
+    f = _make_fn_with_signature(
+        self.__class__.__name__, signature, module, output=False
+    )
+    f = jaxtyped(f, typechecker=typechecker)
+    f(self, **values)
 
 
-def _jaxtyped_typechecker(typechecker):
-    """A decorator added by the import hook to all classes. Only affects dataclasses.
+def _make_fn_with_signature(
+    name: str, signature: inspect.Signature, module: str, output: bool
+):
+    """Dynamically creates a function `fn` with name `name` and signature `signature`.
 
-    Will be called as
-    ```
-    @_jaxtyped_typechecker(beartype.beartype)
-    @dataclasses.dataclass
-    class SomeDataclass:
-        ...
-    ```
+    If `output=True` then `fn` will consume an additional keyword-only argument (in
+    addition to the provided signature), and will directly return this argument. In this
+    case the returned value from `_make_fn_with_signature` is a 2-tuple of `(fn, name)`,
+    where `fn` is the generated function, and `name` is the name of this extra argument.
 
-    After initialisation, this will check that all fields of the dataclass match their
-    specified type annotation.
+    If `output=False` then `fn` will just have a single `pass` statement, and the
+    returned value from `_make_fn_with_signature` will just be `fn`.
+
+    ---
+
+    Note that this function operates by dynamically creating and eval'ing a string, not
+    simply by assigning `__signature__` and `__annotations__`. The latter is enough for
+    typeguard (at least v2), but does not work with beartype (at least v16).
     """
-    # typechecker is expected to probably be either `typeguard.typechecked`, or
-    # `beartype.beartype`, or `None`.
+    pos = []
+    pos_or_key = []
+    varpos = []
+    key = []
+    varkey = []
+    for p in signature.parameters.values():
+        if p.kind == inspect.Parameter.POSITIONAL_ONLY:
+            pos.append(p)
+        elif p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+            pos_or_key.append(p)
+        elif p.kind == inspect.Parameter.VAR_POSITIONAL:
+            varpos.append(p)
+        elif p.kind == inspect.Parameter.KEYWORD_ONLY:
+            key.append(p)
+        elif p.kind == inspect.Parameter.VAR_KEYWORD:
+            varkey.append(p)
+        else:
+            assert False
 
-    if typechecker is None:
-        typechecker = lambda x: x
+    param_names = frozenset(signature.parameters.keys())
+    if output:
+        output_name = _gensym(param_names, prefix="ret")
+        outstr = "return " + output_name
+        param_names = param_names | frozenset({output_name})
+        key.append(inspect.Parameter(output_name, kind=inspect.Parameter.KEYWORD_ONLY))
+    else:
+        outstr = "pass"
 
-    def _wrapper(kls):
-        assert inspect.isclass(kls)
-        if dataclasses.is_dataclass(kls):
-            # This does not check that the arguments passed to `__init__` match the
-            # type annotations. There may be a custom user `__init__`, or a
-            # dataclass-generated `__init__` used alongside
-            # `equinox.field(converter=...)`
+    scope = {name: None}
+    name_to_annotation = {}
+    name_to_default = {}
+    param_triples = (
+        (p.name, p.annotation, p.default) for p in signature.parameters.values()
+    )
+    if output:
+        triples = it.chain(
+            param_triples,
+            [
+                ("return", signature.return_annotation, inspect.Signature.empty),
+                (output_name, Any, inspect.Signature.empty),
+            ],
+        )
+    else:
+        triples = it.chain(
+            param_triples,
+            [("return", signature.return_annotation, inspect.Signature.empty)],
+        )
+    for p_name, p_annotation, p_default in triples:
+        annotation_name = _gensym(frozenset(scope.keys()) | param_names, prefix="T")
+        name_to_annotation[p_name] = annotation_name
+        if p_annotation is inspect.Signature.empty or isinstance(p_annotation, str):
+            # If we have a stringified annotation here it's because the get_type_hints
+            # lookup above failed. Typically this occurs when using a local variable as
+            # the annotation. In this case we really have no idea what the annotation
+            # refers to, so just set it to Any.
+            # This does mean that we don't handle partially-stringified local
+            # annotations, e.g. `type["Foo"]` for some local type `Foo`. Those will
+            # probably just error out. Nothing better we can do about that
+            # unfortunately.
+            scope[annotation_name] = Any
+        else:
+            scope[annotation_name] = p_annotation
+        default_name = _gensym(frozenset(scope.keys()) | param_names, prefix="default")
+        name_to_default[p_name] = default_name
+        scope[default_name] = p_default
 
-            init = kls.__init__
+    argstr_pieces = []
+    if len(pos) > 0:
+        for p in pos:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+        argstr_pieces.append("/")
+    if len(pos_or_key) > 0:
+        for p in pos_or_key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varpos) == 1:
+        [p] = varpos
+        argstr_pieces.append(
+            "*" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varpos) == 0
+        if len(key) > 0:
+            argstr_pieces.append("*")
+    if len(key) > 0:
+        for p in key:
+            argstr_pieces.append(_make_argpiece(p, name_to_annotation, name_to_default))
+    if len(varkey) == 1:
+        [p] = varkey
+        argstr_pieces.append(
+            "**" + _make_argpiece(p, name_to_annotation, name_to_default)
+        )
+    else:
+        assert len(varkey) == 0
+    argstr = ", ".join(argstr_pieces)
 
-            @ft.wraps(init)
-            def __init__(self, *args, **kwargs):
-                init(self, *args, **kwargs)
-                # `kls.__init__` is late-binding to the `__init__` function that we're
-                # in now. (Or to someone else's monkey-patch.) Either way, this checks
-                # that we're in the "top-level" `__init__`, and not one that is being
-                # called via `super()`. We don't want to trigger too early, before all
-                # fields have been assigned.
-                #
-                # We're not checking `if self.__class__ is kls` because Equinox replaces
-                # the with a defrozen version of itself during `__init__`, so the check
-                # wouldn't trigger.
-                #
-                # We're not doing this check by adding it to the end of the metaclass
-                # `__call__`, because Python doesn't allow you monkey-patch metaclasses.
-                if self.__class__.__init__ is kls.__init__:
-                    _check_dataclass_annotations(self, typechecker)
+    if signature.return_annotation is inspect.Signature.empty:
+        retstr = ""
+    else:
+        retstr = f"-> {name_to_annotation['return']}"
 
-            kls.__init__ = __init__
-        return kls
+    fnstr = f"def {name}({argstr}){retstr}:\n    {outstr}"
+    exec(fnstr, scope)
+    fn = scope[name]
+    fn.__module__ = module
+    assert fn is not None
+    if output:
+        return fn, output_name
+    else:
+        return fn
 
-    return _wrapper
+
+def _gensym(names: frozenset[str], prefix: str) -> str:
+    assert prefix.isidentifier()
+    output_index = 0
+    output_name = prefix + str(output_index)
+    while output_name in names:
+        output_index += 1
+        output_name = prefix + str(output_index)
+    assert output_name.isidentifier()
+    return output_name
+
+
+def _make_argpiece(p, name_to_annotation, name_to_default):
+    if p.default is inspect.Signature.empty:
+        return f"{p.name}: {name_to_annotation[p.name]}"
+    else:
+        return f"{p.name}: {name_to_annotation[p.name]} = {name_to_default[p.name]}"
+
+
+def _get_problem_arg(
+    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.
+    """
+    # No performance concerns, as this is only used when we're about to raise an error
+    # 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
+        )
+        fn = typechecker(fn)  # but no `jaxtyped`; keep the same environment.
+        try:
+            fn(*args, **kwargs)
+        except Exception:
+            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.
+        return ""
+
+
+def _remove_typing(x):
+    x = str(x)
+    x = x.replace(" jaxtyping.", " ")
+    x = x.replace("[jaxtyping.", "[")
+    x = x.replace("'jaxtyping.", "'")
+    x = x.replace(" typing.", " ")
+    x = x.replace("[typing.", "[")
+    x = x.replace("'typing.", "'")
+    return x
+
+
+def _pformat(x, short_self: bool):
+    # No performance concerns from delayed imports -- this is only used when we're about
+    # to raise an error anyway.
+    try:
+        # TODO(kidger): this is pretty ugly. We have a circular dependency
+        # equinox->jaxtyping->equinox. We could consider moving all the pretty-printing
+        # code from equinox into jaxtyping maybe? Or into some shared dependency?
+        import equinox as eqx
+
+        pformat = eqx.tree_pformat
+        if short_self:
+            try:
+                self = x["self"]
+            except KeyError:
+                pass
+            else:
+                is_self = lambda y: y is self
+                pformat = ft.partial(pformat, truncate_leaf=is_self)
+    except Exception:
+        import pprint
+
+        pformat = ft.partial(pprint.pformat, indent=2, compact=True)
+    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."""
+
+
+def _no_jaxtyping_note(e: Exception) -> bool:
+    """Checks if any of the exception's notes are from jaxtyping."""
+    try:
+        notes = e.__notes__
+    except AttributeError:
+        return True
+    else:
+        for note in notes:
+            if isinstance(note, _jaxtyping_note_str):
+                return False
+        return True
+
+
+_spacer = "--------------------\n"

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

@@ -78,18 +78,13 @@ def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None)
     #    changing the typechecker will hit a different cache.
     # Version 7: Using the same md5 hash of the `typechecker` argument
     #    for importlib and decorator lookup.
+    # Version 8: Now using new-style `jaxtyped(typechecker=...)` rather than old-style
+    #    double-decorators.
     return cache_from_source(
-        path, debug_override, optimization=f"jaxtyping7{typechecker_hash}"
+        path, debug_override, optimization=f"jaxtyping8{typechecker_hash}"
     )
 
 
-def _dot_lookup(*elements):
-    out = ast.Name(id=elements[0], ctx=ast.Load())
-    for element in elements[1:]:
-        out = ast.Attribute(out, element, ctx=ast.Load())
-    return out
-
-
 class Typechecker:
     lookup = {}
 
@@ -114,7 +109,7 @@ class Typechecker:
 
         elif typechecker is None:
             # If it is None, ignore it silently (use dummy decorator)
-            self.hash = 0
+            self.hash = "0"
             Typechecker.lookup[self.hash] = lambda x, *_, **__: x
         else:
             # Passed typechecker is invalid
@@ -130,7 +125,7 @@ class Typechecker:
         if self.ast is None:
             self.ast = (
                 ast.parse(
-                    f"@jaxtyping._import_hook.Typechecker.lookup['{self.hash}']\n"
+                    f"@jaxtyping.jaxtyped(typechecker=jaxtyping._import_hook.Typechecker.lookup['{self.hash}'])\n"
                     "def _():\n    ..."
                 )
                 .body[0]
@@ -162,10 +157,9 @@ class JaxtypingTransformer(ast.NodeVisitor):
         return node
 
     def visit_ClassDef(self, node: ast.ClassDef):
-        func = _dot_lookup("jaxtyping", "_decorator", "_jaxtyped_typechecker")
-        node.decorator_list.insert(
-            0, ast.Call(func, [self._typechecker.get_ast()], keywords=[])
-        )
+        # Place at the start of the decorator list, so that `@dataclass` decorators get
+        # called first.
+        node.decorator_list.insert(0, self._typechecker.get_ast())
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -175,8 +169,13 @@ class JaxtypingTransformer(ast.NodeVisitor):
         has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
         has_annotated_return = bool(node.returns)
         if has_annotated_args or has_annotated_return:
-            # Place at the end of the decorator list, as otherwise we wrap e.g.
-            # `jax.custom_{jvp,vjp}` and lose the ability to `defjvp` etc.
+            # Place at the end of the decorator list, because:
+            # - as otherwise we wrap e.g. `jax.custom_{jvp,vjp}` and lose the ability
+            #     to `defjvp` etc.
+            # - decorators frequently remove annotations from functions, and we'd like
+            #     to use those annotations.
+            # - typeguard in particular wants to be at the end of the decorator list, as
+            #     it works by recompling the wrapped function.
             #
             # Note that the counter-argument here is that we'd like to place this
             # at the start of the decorator list, in case a typechecking annotation
@@ -184,14 +183,6 @@ class JaxtypingTransformer(ast.NodeVisitor):
             # case we're just going to have to need to ask the user to remove their
             # typechecking annotation (and let this decorator do it instead).
             # It's more important we be compatible with normal JAX code.
-            #
-            #
-            # FWIW, typeguard also wants to be at the end of the decorator list, as it
-            # works by recompiling the wrapped function.
-            node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
-            # Place typechecker at the end of the decorator list, as decorators
-            # frequently remove annotations from functions and we'd like to
-            # use those annotations.
             node.decorator_list.append(self._typechecker.get_ast())
 
         self._parents.append(node)
@@ -290,7 +281,8 @@ class ImportHookManager:
 # Deliberately no default for `typechecker` so that folks must opt-in to not having
 # a typechecker.
 def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optional[str]):
-    """Automatically apply `@jaxtyped`, and optionally a type checker, as decorators.
+    """Automatically apply the `@jaxtyped(typechecker=typechecker)` decorator to every
+    function and dataclass over a whole codebase.
 
     !!! Tip "Usage"
 
@@ -298,18 +290,19 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
         from jaxtyping import install_import_hook
         # Plus any one of the following:
 
-        # decorate @jaxtyped and @typeguard.typechecked
+        # decorate `@jaxtyped(typechecker=typeguard.typechecked)`
         with install_import_hook("foo", "typeguard.typechecked"):
             import foo          # Any module imported inside this `with` block, whose
             import foo.bar      # name begins with the specified string, will
             import foo.bar.qux  # automatically have both `@jaxtyped` and the specified
-                                # typechecker applied to all of their functions.
+                                # typechecker applied to all of their functions and
+                                # dataclasses.
 
-        # decorate @jaxtyped and @beartype.beartype
+        # decorate `@jaxtyped(typechecker=beartype.beartype)`
         with install_import_hook("foo", "beartype.beartype"):
             ...
 
-        # decorate only @jaxtyped (if you want that for some reason)
+        # decorate only `@jaxtyped` (if you want that for some reason)
         with install_import_hook("foo", None):
             ...
         ```
@@ -326,22 +319,9 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
         install_import_hook(["foo", "bar.baz"], ...)
         ```
 
-    The import hook will automatically decorate all functions, and check the attributes
-    assigned to dataclasses.
-
-    If the function already has any decorators on it, then both the `@jaxtyped` and the
-    typechecker decorators will get added at the bottom of the decorator list, e.g.
-    ```python
-    @some_other_decorator
-    @jaxtyped
-    @beartype.beartype
-    def foo(...): ...
-    ```
-
     **Arguments:**:
 
-    - `modules`: the names of the modules in which to automatically apply `@jaxtyped`
-        and `@typechecked`.
+    - `modules`: the names of the modules in which to automatically apply `@jaxtyped`.
     - `typechecker`: the module and function of the typechecker you want to use, as a
         string. For example `typechecker="typeguard.typechecked"`, or
         `typechecker="beartype.beartype"`. You may pass `typechecker=None` if you do not
@@ -351,7 +331,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
 
     A context manager that uninstalls the hook on exit, or when you call `.uninstall()`.
 
-    ??? Example "Example: end-user script"
+    !!! Example "Example: end-user script"
 
         ```python
         ### entry_point.py
@@ -366,7 +346,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
             ...
         ```
 
-    ??? Example "Example: writing a library"
+    !!! Example "Example: writing a library"
 
         ```python
         ### __init__.py
@@ -378,30 +358,6 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
                                                  # so will be hook'd.
         ```
 
-    ??? info "Pytest hook"
-
-        The import hook can be installed at test-time only, as a pytest hook. From the
-        command line the syntax is:
-        ```
-        pytest --jaxtyping-packages=foo,bar.baz,beartype.beartype
-        ```
-        or in `pyproject.toml`:
-        ```toml
-        [tool.pytest.ini_options]
-        addopts = "--jaxtyping-packages=foo,bar.baz,beartype.beartype"
-        ```
-        or in `pytest.ini`:
-        ```ini
-        [pytest]
-        addopts = --jaxtyping-packages=foo,bar.baz,beartype.beartype
-        ```
-        This example will apply the import hook to all modules whose names start with
-        either `foo` or `bar.baz`. The typechecker used in this example is
-        `beartype.beartype`.
-
-        (This is the author's preferred approach to performing runtime type-checking
-        with jaxtyping!)
-
     !!! warning
 
         Stringified dataclass annotations, e.g.
@@ -423,6 +379,29 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
             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
+        at the bottom of the decorator list, e.g.
+        ```python
+        @some_other_decorator
+        @jaxtyped(typechecker=beartype.beartype)
+        def foo(...): ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = jax.custom_jvp` etc.
+
+        If a class already has any decorators in it, then `@jaxtyped` will get added to
+        the top of the decorator list, e.g.
+        ```python
+        @jaxtyped(typechecker=beartype.beartype)
+        @some_other_decorator
+        class A:
+            ...
+        ```
+        This is to support the common case in which
+        `some_other_decorator = dataclasses.dataclass`.
     """  # noqa: E501
 
     if isinstance(modules, str):

jaxtyping/_indirection.py

@@ -39,6 +39,7 @@ from typing import (
     Annotated as Integer,  # noqa: F401
     Annotated as Key,  # noqa: F401
     Annotated as Num,  # noqa: F401
+    Annotated as Real,  # noqa: F401
     Annotated as Shaped,  # noqa: F401
     Annotated as UInt,  # noqa: F401
     Annotated as UInt8,  # noqa: F401
@@ -47,5 +48,8 @@ from typing import (
     Annotated as UInt64,  # noqa: F401
 )
 
-from jax import Array as Scalar  # noqa: F401
+from jax import (
+    Array as PRNGKeyArray,  # noqa: F401
+    Array as Scalar,  # noqa: F401
+)
 from jax.typing import ArrayLike as ScalarLike  # noqa: F401

jaxtyping/_pytree_type.py

@@ -19,11 +19,21 @@
 
 import functools as ft
 import typing
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 import jax.tree_util as jtu
 import typeguard
 
+from ._errors import AnnotationError
+from ._storage import (
+    clear_treeflatten_memo,
+    clear_treepath_memo,
+    get_shape_memo,
+    set_shape_memo,
+    set_treeflatten_memo,
+    set_treepath_memo,
+)
+
 
 _T = TypeVar("_T")
 
@@ -44,27 +54,133 @@ class _MetaPyTree(type):
     def __instancecheck__(cls, obj):
         if not hasattr(cls, "leaftype"):
             return True  # Just `isinstance(x, PyTree)`
+        # Handle beartype doing `isinstance(None, hint)` to check if
+        # is `instance`able.
+        if obj is None:
+            return True
 
-        # We could use `isinstance` here but that would fail for more complicated
-        # types, e.g. PyTree[Tuple[int]]. So at least internally we make a particular
-        # choice of typechecker.
-        #
-        # Deliberately not using @jaxtyped so that we share the same `memo` as whatever
-        # dynamic context we're currently in.
-        @typeguard.typechecked
-        def accepts_leaftype(x: cls.leaftype):
-            pass
+        single_memo, variadic_memo, pytree_memo, arg_memo = get_shape_memo()
+        single_memo_bak = single_memo.copy()
+        variadic_memo_bak = variadic_memo.copy()
+        pytree_memo_bak = pytree_memo.copy()
+        arg_memo_bak = arg_memo.copy()
+        try:
+            out = cls._check(obj, pytree_memo)
+        except Exception:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            raise
+        if out:
+            return True
+        else:
+            set_shape_memo(
+                single_memo_bak, variadic_memo_bak, pytree_memo_bak, arg_memo_bak
+            )
+            return False
 
-        def is_leaftype(x):
-            try:
-                accepts_leaftype(x)
-            except _TypeCheckError:
+    def _check(cls, obj, pytree_memo):
+        if cls.leaftype is Any:
+
+            def is_flatten_leaftype(x):
                 return False
-            else:
+
+            def is_check_leaftype(x):
                 return True
 
-        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
-        return all(map(is_leaftype, leaves))
+        else:
+            # We could use `isinstance` here but that would fail for more complicated
+            # types, e.g. PyTree[tuple[int]]. So at least internally we make a
+            # particular choice of typechecker.
+            #
+            # Deliberately not using @jaxtyped so that we share the same `memo` as
+            # whatever dynamic context we're currently in.
+            @typeguard.typechecked
+            def accepts_leaftype(x: cls.leaftype):
+                pass
+
+            def is_leaftype(x):
+                try:
+                    accepts_leaftype(x)
+                except _TypeCheckError:
+                    return False
+                else:
+                    return True
+
+            is_flatten_leaftype = is_check_leaftype = is_leaftype
+
+        set_treeflatten_memo()
+        try:
+            leaves, structure = jtu.tree_flatten(obj, is_leaf=is_flatten_leaftype)
+        finally:
+            clear_treeflatten_memo()
+        if cls.structure is not None:
+            if cls.structure.isidentifier():
+                try:
+                    prev_structure = pytree_memo[cls.structure]
+                except KeyError:
+                    pytree_memo[cls.structure] = structure
+                else:
+                    if prev_structure != structure:
+                        return False
+            else:
+                named_pytree = 0
+                pieces = cls.structure.split()
+                if pieces[0] == "...":
+                    pieces = pieces[1:]
+                    prefix = False
+                    suffix = True
+                elif pieces[-1] == "...":
+                    pieces = pieces[:-1]
+                    prefix = True
+                    suffix = False
+                else:
+                    prefix = False
+                    suffix = False
+                for identifier in pieces:
+                    try:
+                        prev_structure = pytree_memo[identifier]
+                    except KeyError as e:
+                        raise AnnotationError(
+                            f"Cannot process composite structure '{cls.structure}' "
+                            f"as the structure name {identifier} has not been seen "
+                            "before."
+                        ) from e
+                    # Not using `PyTreeDef.compose` due to JAX bug #18218.
+                    prev_pytree = jtu.tree_unflatten(
+                        prev_structure, [0] * prev_structure.num_leaves
+                    )
+                    named_pytree = jtu.tree_map(lambda _: prev_pytree, named_pytree)
+                named_structure = jtu.tree_structure(named_pytree)
+                if prefix:
+                    dummy_pytree = jtu.tree_unflatten(structure, [0] * len(leaves))
+                    dummy_named = jtu.tree_unflatten(
+                        named_structure, [0] * named_structure.num_leaves
+                    )
+                    try:
+                        jtu.tree_map(lambda _, __: 0, dummy_named, dummy_pytree)
+                    except ValueError:
+                        return False
+                elif suffix:
+                    has_structure = lambda x: jtu.tree_structure(x) == named_structure
+                    dummy_pytree = jtu.tree_unflatten(structure, [0] * len(leaves))
+                    dummy_leaves = jtu.tree_leaves(dummy_pytree, is_leaf=has_structure)
+                    if any(not has_structure(x) for x in dummy_leaves):
+                        return False
+                else:
+                    if structure != named_structure:
+                        return False
+
+        try:
+            for leaf_index, leaf in enumerate(leaves):
+                if cls.structure is not None:
+                    set_treepath_memo(leaf_index, cls.structure)
+                if not is_check_leaftype(leaf):
+                    return False
+                clear_treepath_memo()
+        finally:
+            clear_treepath_memo()
+        return True
 
     # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
     # the custom __instancecheck__ that we want.
@@ -75,10 +191,54 @@ class _MetaPyTree(type):
     # has __module__ "types", e.g. we get types.PyTree[int].
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
-        name = str(_FakePyTree[item])
+        if isinstance(item, tuple):
+            if len(item) == 2:
 
-        class X(PyTree):
-            leaftype = item
+                class X(PyTree):
+                    leaftype = item[0]
+                    structure = item[1].strip()
+
+                if not isinstance(X.structure, str):
+                    raise ValueError(
+                        "The structure annotation `struct` in "
+                        "`jaxtyping.PyTree[leaftype, struct]` must be be a string, "
+                        f"e.g. `jaxtyping.PyTree[leaftype, 'T']`. Got '{X.structure}'."
+                    )
+                pieces = X.structure.split()
+                if len(pieces) == 0:
+                    raise ValueError(
+                        "The string `struct` in `jaxtyping.PyTree[leaftype, struct]` "
+                        "cannot be the empty string."
+                    )
+                for piece_index, piece in enumerate(pieces):
+                    if (piece_index == 0) or (piece_index == len(pieces) - 1):
+                        if piece == "...":
+                            continue
+                    if not piece.isidentifier():
+                        raise ValueError(
+                            "The string `struct` in "
+                            "`jaxtyping.PyTree[leaftype, struct]` must be be a "
+                            "whitespace-separated sequence of identifiers, e.g. "
+                            "`jaxtyping.PyTree[leaftype, 'T']` or "
+                            "`jaxtyping.PyTree[leaftype, 'foo bar']`.\n"
+                            "(Here, 'identifier' is used in the same sense as in "
+                            "regular Python, i.e. a valid variable name.)\n"
+                            f"Got piece '{piece}' in overall structure '{X.structure}'."
+                        )
+                name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
+            else:
+                raise ValueError(
+                    "The subscript `foo` in `jaxtyping.PyTree[foo]` must either be a "
+                    "leaf type, e.g. `PyTree[int]`, or a 2-tuple of leaf and "
+                    "structure, e.g. `PyTree[int, 'T']`. Received a tuple of length "
+                    f"{len(item)}."
+                )
+        else:
+            name = str(_FakePyTree[item])
+
+            class X(PyTree):
+                leaftype = item
+                structure = None
 
         X.__name__ = name
         X.__qualname__ = name
@@ -107,10 +267,56 @@ else:
     PyTree.__module__ = "jaxtyping"
 PyTree.__doc__ = """Represents a PyTree.
 
-Each PyTree is denoted by a type `PyTree[LeafType]`, such as `PyTree[int]` or
-`PyTree[Union[str, Float32[Array, "b c"]]]`.
+Annotations of the following sorts are supported:
+```python
+a: PyTree
+b: PyTree[LeafType]
+c: PyTree[LeafType, "T"]
+d: PyTree[LeafType, "S T"]
+e: PyTree[LeafType, "... T"]
+f: PyTree[LeafType, "T ..."]
+```
 
-You can leave off the `[...]`, in which case `PyTree` is simply a suggestively-named
-alternative to `Any`.
-([By definition all types are PyTrees.](https://jax.readthedocs.io/en/latest/pytrees.html))
+These correspond to:
+
+a. A plain `PyTree` can be used an annotation, in which case `PyTree` is simply a
+    suggestively-named alternative to `Any`.
+    ([By definition all types are PyTrees.](https://jax.readthedocs.io/en/latest/pytrees.html))
+
+b. `PyTree[LeafType]` denotes a PyTree all of whose leaves match `LeafType`. For
+    example, `PyTree[int]` or `PyTree[Union[str, Float32[Array, "b c"]]]`.
+
+c. A structure name can also be passed. In this case
+    `jax.tree_util.tree_structure(...)` will be called, and bound to the structure name.
+    This can be used to mark that multiple PyTrees all have the same structure:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "T"]):
+        ...
+    ```
+    Structures are bound to names in the same way as array shape annotations, i.e.
+    within the thread-local dynamic context of a [`jaxtyping.jaxtyped`][] decorator.
+
+d. A composite structure can be declared. In this case the variable must have a PyTree
+    structure each to the composition of multiple previously-bound PyTree structures.
+    For example:
+    ```python
+    def f(x: PyTree[int, "T"], y: PyTree[int, "S"], z: PyTree[int, "S T"]):
+        ...
+
+    x = (1, 2)
+    y = {"key": 3}
+    z = {"key": (4, 5)}  # structure is the composition of the structures of `y` and `z`
+    f(x, y, z)
+    ```
+    When performing runtime type-checking, all the individual pieces must have already
+    been bound to structures, otherwise the composite structure check will throw an error.
+
+e. A structure can begin with a `...`, to denote that the lower levels of the PyTree
+    must match the declared structure, but the upper levels can be arbitrary. As in the
+    previous case, all named pieces must already have been seen and their structures
+    bound.
+
+f. A structure can end with a `...`, to denote that the PyTree must be a prefix of the
+    declared structure, but the lower levels can be arbitrary. As in the previous two
+    cases, all named pieces must already have been seen and their structures bound.
 """  # noqa: E501

jaxtyping/_storage.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2022 Google LLC
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# 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 threading
+from typing import Any, Optional
+
+from ._errors import AnnotationError
+
+
+_shape_storage = threading.local()
+
+
+def _has_shape_memo():
+    return hasattr(_shape_storage, "memo_stack") and len(_shape_storage.memo_stack) != 0
+
+
+def get_shape_memo():
+    if _has_shape_memo():
+        single_memo, variadic_memo, pytree_memo, arguments = _shape_storage.memo_stack[
+            -1
+        ]
+    else:
+        # `isinstance` happening outside any @jaxtyped decorators, e.g. at the
+        # global scope. In this case just create a temporary memo, since we're not
+        # going to be comparing against any stored values anyway.
+        single_memo = {}
+        variadic_memo = {}
+        pytree_memo = {}
+        arguments = {}
+    return single_memo, variadic_memo, pytree_memo, arguments
+
+
+def set_shape_memo(single_memo, variadic_memo, pytree_memo, arg_memo) -> None:
+    if _has_shape_memo():
+        _shape_storage.memo_stack[-1] = (
+            single_memo,
+            variadic_memo,
+            pytree_memo,
+            arg_memo,
+        )
+
+
+def push_shape_memo(arguments: dict[str, Any]):
+    try:
+        memo_stack = _shape_storage.memo_stack
+    except AttributeError:
+        # Can't be done when `_stack_storage` is created for reasons I forget.
+        memo_stack = _shape_storage.memo_stack = []
+    memos = ({}, {}, {}, arguments.copy())
+    memo_stack.append(memos)
+    return memos
+
+
+def pop_shape_memo() -> None:
+    _shape_storage.memo_stack.pop()
+
+
+_treepath_storage = threading.local()
+
+
+def clear_treepath_memo() -> None:
+    _treepath_storage.value = None
+
+
+def set_treepath_memo(index: Optional[int], structure: str) -> None:
+    if hasattr(_treepath_storage, "value") and _treepath_storage.value is not None:
+        raise AnnotationError(
+            "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:
+        # Appears in error messages, so human-readable
+        _treepath_storage.value = f"(Leaf {index} in structure {structure}) "
+
+
+def get_treepath_memo() -> str:
+    if not hasattr(_treepath_storage, "value") or _treepath_storage.value is None:
+        raise AnnotationError(
+            "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
+
+
+_treeflatten_storage = threading.local()
+
+
+def clear_treeflatten_memo() -> None:
+    _treeflatten_storage.value = False
+
+
+def set_treeflatten_memo():
+    _treeflatten_storage.value = True
+
+
+def get_treeflatten_memo():
+    try:
+        return _treeflatten_storage.value
+    except AttributeError:
+        return False

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.23"
+version = "0.2.25"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"

test/conftest.py

@@ -41,6 +41,30 @@ def typecheck(request):
     return request.param
 
 
+@pytest.fixture(params=(False, True))
+def jaxtyp(request):
+    import jaxtyping
+
+    if request.param:
+        # New-style
+        # @jaxtyping.jaxtyped(typechecker=typechecker)
+        # def f(...)
+        return lambda typechecker: jaxtyping.jaxtyped(typechecker=typechecker)
+    else:
+        # Old-style
+        # @jaxtyping.jaxtyped
+        # @typechecker
+        # 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
+
+        return impl
+
+
 @pytest.fixture()
 def getkey():
     def _getkey():

test/test_array.py

@@ -17,6 +17,7 @@
 # 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 dataclasses as dc
 import sys
 from typing import get_args, get_origin, Union
 
@@ -28,21 +29,24 @@ import torch
 
 from jaxtyping import (
     AbstractDtype,
+    AnnotationError,
     Array,
     ArrayLike,
+    Bool,
     Float,
     Float32,
     jaxtyped,
+    Key,
     PRNGKeyArray,
+    Scalar,
     Shaped,
 )
 
 from .helpers import ParamError, ReturnError
 
 
-def test_basic(typecheck):
-    @jaxtyped
-    @typecheck
+def test_basic(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "..."]):
         pass
 
@@ -81,16 +85,14 @@ def test_dtypes():
             assert key == val.__name__
 
 
-def test_return(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_return(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
         return jnp.transpose(x)
 
     g(jr.normal(getkey(), (3, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
         return jnp.transpose(x)
 
@@ -98,9 +100,8 @@ def test_return(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)))
 
 
-def test_two_args(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_two_args(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
         return x @ y
 
@@ -108,8 +109,7 @@ def test_two_args(typecheck, getkey):
     with pytest.raises(ParamError):
         g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
         return x @ y
 
@@ -118,9 +118,8 @@ def test_two_args(typecheck, getkey):
         h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (5, 4)))
 
 
-def test_any_dtype(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_any_dtype(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
         return x
 
@@ -135,14 +134,12 @@ def test_any_dtype(typecheck, getkey):
         g(jr.normal(getkey(), (1,)))
 
 
-def test_nested_jaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_jaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"], transpose: bool) -> Float32[Array, "c b"]:
         return h(x, transpose)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, "c b"], transpose: bool) -> Float32[Array, "b c"]:
         if transpose:
             return jnp.transpose(x)
@@ -156,9 +153,8 @@ def test_nested_jaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)), False)
 
 
-def test_nested_nojaxtyped(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_nested_nojaxtyped(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]):
         return h(x)
 
@@ -170,9 +166,8 @@ def test_nested_nojaxtyped(typecheck, getkey):
         g(jr.normal(getkey(), (2, 3)))
 
 
-def test_isinstance(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_isinstance(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "b c"]) -> Float32[Array, " z"]:
         y = jnp.transpose(x)
         assert isinstance(y, Float32[Array, "c b"])
@@ -186,9 +181,8 @@ def test_isinstance(typecheck, getkey):
     g(jr.normal(getkey(), (2, 3)))
 
 
-def test_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "4 5 foo"], y: Float32[Array, " foo"]
     ) -> Float32[Array, "4 5"]:
@@ -203,9 +197,8 @@ def test_fixed(typecheck, getkey):
         g(c, b)
 
 
-def test_anonymous(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "foo _"], y: Float32[Array, " _"]):
         pass
 
@@ -214,9 +207,8 @@ def test_anonymous(typecheck, getkey):
     g(a, b)
 
 
-def test_named_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_named_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(
         x: Float32[Array, "*batch foo"],
         y: Float32[Array, " *batch"],
@@ -239,8 +231,7 @@ def test_named_variadic(typecheck, getkey):
     with pytest.raises(ParamError):
         g(a2, b1, c)
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def h(x: Float32[Array, " foo *batch"], y: Float32[Array, " foo *batch bar"]):
         pass
 
@@ -254,9 +245,8 @@ def test_named_variadic(typecheck, getkey):
         h(b, c)
 
 
-def test_anonymous_variadic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_anonymous_variadic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "... foo"], y: Float32[Array, " foo"]):
         pass
 
@@ -276,9 +266,8 @@ def test_anonymous_variadic(typecheck, getkey):
         g(a3, c)
 
 
-def test_broadcast_fixed(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_fixed(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, "#4"]):
         pass
 
@@ -289,9 +278,8 @@ def test_broadcast_fixed(typecheck, getkey):
         g(jr.normal(getkey(), (3,)))
 
 
-def test_broadcast_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " #foo"], y: Float32[Array, " #foo"]):
         pass
 
@@ -313,9 +301,8 @@ def test_broadcast_named(typecheck, getkey):
         g(b, a)
 
 
-def test_broadcast_variadic_named(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_broadcast_variadic_named(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: Float32[Array, " *#foo"], y: Float32[Array, " *#foo"]):
         pass
 
@@ -372,29 +359,71 @@ def test_broadcast_variadic_named(typecheck, getkey):
         g(o, a)
 
 
+def test_variadic_mixed_broadcast(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: Float[Array, " *foo"], y: Float[Array, " #*foo"]):
+        pass
+
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with pytest.raises(ParamError):
+        f(a, b)
+
+    c = jr.normal(getkey(), (7, 3, 2))
+    d = jr.normal(getkey(), (1, 2))
+    f(c, d)
+
+
+def test_variadic_mixed_broadcast2(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: Float[Array, " *#foo"], y: Float[Array, " *foo"]):
+        pass
+
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with pytest.raises(ParamError):
+        f(a, b)
+
+    c = jr.normal(getkey(), (1, 2))
+    d = jr.normal(getkey(), (7, 3, 2))
+    f(c, d)
+
+
+def test_variadic_mixed_broadcast3(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(
+        x: Float[Array, "*B L D"],
+        *,
+        y: Float[Array, "*#B J d"],
+        z: Bool[Array, "*B L J"],
+    ) -> Float[Array, "*B L D"]:
+        return x
+
+    x = jr.normal(getkey(), (2, 7, 3, 2, 2))
+    y = jr.bernoulli(getkey(), shape=(2, 7, 3, 2, 2))
+    z = jr.normal(getkey(), (2, 7, 1, 2, 2))
+    f(x, y=z, z=y)
+
+
 def test_no_commas():
     with pytest.raises(ValueError):
         Float32[Array, "foo, bar"]
 
 
-def test_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x[1:]
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x])
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_make_slice(x: Float32[Array, " dim"]) -> Float32[Array, " dim-1"]:
         return x
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def bad_cat(x: Float32[Array, " dim"]) -> Float32[Array, " 2*dim"]:
         return jnp.concatenate([x, x, x])
 
@@ -414,17 +443,64 @@ def test_symbolic(typecheck, getkey):
         bad_cat(x)
 
 
-def test_incomplete_symbolic(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def foo(x: Float32[Array, " 2*dim"]):
         pass
 
     x = jr.normal(getkey(), (4,))
-    with pytest.raises(NameError):
+    with pytest.raises(AnnotationError):
         foo(x)
 
 
+def test_deferred_symbolic_good(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim}"]:
+        return jnp.full((dim,), fill)
+
+    class A:
+        size = 5
+
+        @jaxtyp(typecheck)
+        def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}"]:
+            return jnp.full((self.size,), fill)
+
+    foo(3, jnp.array(0.0))
+    A().bar(jnp.array(0.0))
+
+
+def test_deferred_symbolic_bad(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def foo(dim: int, fill: Float[Array, ""]) -> Float[Array, " {dim-1}"]:
+        return jnp.full((dim,), fill)
+
+    class A:
+        size = 5
+
+        @jaxtyp(typecheck)
+        def bar(self, fill: Float[Array, ""]) -> Float[Array, " {self.size}-1"]:
+            return jnp.full((self.size,), fill)
+
+    with pytest.raises(ReturnError):
+        foo(3, jnp.array(0.0))
+
+    with pytest.raises(ReturnError):
+        A().bar(jnp.array(0.0))
+
+
+def test_deferred_symbolic_dataclass(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    @dc.dataclass
+    class A:
+        value: int
+        array: Float[Array, " {self.value}"]
+
+    A(3, jnp.zeros(3))
+
+    with pytest.raises(ParamError):
+        A(3, jnp.zeros(4))
+
+
 def test_arraylike(typecheck, getkey):
     floatlike1 = Float32[ArrayLike, ""]
     floatlike2 = Float[ArrayLike, ""]
@@ -436,22 +512,18 @@ def test_arraylike(typecheck, getkey):
     assert set(get_args(floatlike1)) == {
         Float32[Array, ""],
         Float32[np.ndarray, ""],
-        Float32[np.bool_, ""],
         Float32[np.number, ""],
         float,
     }
     assert set(get_args(floatlike2)) == {
         Float[Array, ""],
         Float[np.ndarray, ""],
-        Float[np.bool_, ""],
         Float[np.number, ""],
         float,
     }
     assert set(get_args(floatlike3)) == {
         Float32[Array, "4"],
         Float32[np.ndarray, "4"],
-        Float32[np.bool_, "4"],
-        Float32[np.number, "4"],
     }
 
     shaped1 = Shaped[ArrayLike, ""]
@@ -471,8 +543,6 @@ def test_arraylike(typecheck, getkey):
     assert set(get_args(shaped2)) == {
         Shaped[Array, "4"],
         Shaped[np.ndarray, "4"],
-        Shaped[np.bool_, "4"],
-        Shaped[np.number, "4"],
     }
 
 
@@ -523,14 +593,13 @@ def test_py310_unions():
     assert isinstance(x, get_args(y))
 
 
-def test_key(typecheck):
-    @jaxtyped
-    @typecheck
+def test_key(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def f(x: PRNGKeyArray):
         pass
 
-    x = jr.PRNGKey(0)
-    f(x)
+    f(jr.key(0))
+    f(jr.PRNGKey(0))
 
     with pytest.raises(ParamError):
         f(object())
@@ -542,7 +611,31 @@ def test_key(typecheck):
         f(jnp.array(3.0))
 
 
-def test_extension(typecheck, getkey):
+def test_key_dtype(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f1(x: Key[Array, ""]):
+        pass
+
+    @jaxtyp(typecheck)
+    def f2(x: Key[Scalar, ""]):
+        pass
+
+    for f in (f1, f2):
+        f(jr.key(0))
+
+        with pytest.raises(ParamError):
+            f(jr.PRNGKey(0))
+        with pytest.raises(ParamError):
+            f(object())
+        with pytest.raises(ParamError):
+            f(1)
+        with pytest.raises(ParamError):
+            f(jnp.array(3))
+        with pytest.raises(ParamError):
+            f(jnp.array(3.0))
+
+
+def test_extension(jaxtyp, typecheck, getkey):
     X = Shaped[Array, "a b"]
     Y = Shaped[X, "c d"]
     Z = Shaped[Array, "c d a b"]
@@ -551,8 +644,7 @@ def test_extension(typecheck, getkey):
     X = Float[Array, "a"]
     Y = Float[X, "b"]
 
-    @jaxtyped
-    @typecheck
+    @jaxtyp(typecheck)
     def f(a: X, b: Y):
         ...
 
@@ -575,3 +667,17 @@ def test_extension(typecheck, getkey):
     g(jr.split(jr.PRNGKey(0)))
     with pytest.raises(ParamError):
         g(jr.split(jr.PRNGKey(0), 3))
+
+
+def test_scalar_variadic_dim():
+    assert Float[float, "..."] is float
+    assert Float[float, "#*shape"] is float
+
+    # This one is a bit weird -- it should really also assert that shape==(), but we
+    # don't implement that.
+    assert Float[float, "*shape"] is float
+
+
+def test_scalar_dtype_mismatch():
+    with pytest.raises(ValueError):
+        Float[bool, "..."]

test/test_decorator.py

@@ -1,52 +1,57 @@
 import abc
 
-from jaxtyping import jaxtyped
+import jax.random as jr
+import pytest
+
+from jaxtyping import Array, Float, jaxtyped
+
+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
 
@@ -75,3 +80,89 @@ def test_abstractmethod():
 def test_property():
     assert N().j1 == 3
     assert N().j2 == 4
+
+
+def test_context(getkey):
+    a = jr.normal(getkey(), (3, 4))
+    b = jr.normal(getkey(), (5,))
+    with jaxtyped("context"):
+        assert isinstance(a, Float[Array, "foo bar"])
+        assert not isinstance(b, Float[Array, "foo"])
+    assert isinstance(a, Float[Array, "foo bar"])
+    assert isinstance(b, Float[Array, "foo"])
+
+
+def test_varargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(*args):
+        pass
+
+    f(1, 2)
+
+
+def test_varkwargs(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(**kwargs):
+        pass
+
+    f(a=1, b=2)
+
+
+def test_defaults(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: int, y=1):
+        pass
+
+    f(1)
+
+
+class _GlobalFoo:
+    pass
+
+
+def test_global_stringified_annotation(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: "_GlobalFoo") -> "_GlobalFoo":
+        return x
+
+    f(_GlobalFoo())
+
+    @jaxtyp(typecheck)
+    def g(x: int) -> "_GlobalFoo":
+        return x
+
+    @jaxtyp(typecheck)
+    def h(x: "_GlobalFoo") -> int:
+        return x
+
+    with pytest.raises(ReturnError):
+        g(1)
+
+    with pytest.raises(ParamError):
+        h(1)
+
+
+# This test does not use `jaxtyp(typecheck)` because typeguard does some evil stack
+# frame introspection to try and grab local variables.
+def test_local_stringified_annotation(typecheck):
+    class LocalFoo:
+        pass
+
+    @jaxtyped(typechecker=typecheck)
+    def f(x: "LocalFoo") -> "LocalFoo":
+        return x
+
+    f(LocalFoo())
+
+    with pytest.warns(match="As of jaxtyping version 0.2.24"):
+
+        @jaxtyped
+        @typecheck
+        def g(x: "LocalFoo") -> "LocalFoo":
+            return x
+
+    g(LocalFoo())
+
+    # 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.

test/test_import_hook.py

@@ -29,7 +29,7 @@ import pytest
 import jaxtyping
 
 
-_here = pathlib.Path(__file__).resolve().parent
+_here = pathlib.Path(__file__).parent
 
 
 try:

test/test_messages.py

@@ -0,0 +1,104 @@
+from typing import Any
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Array, Float, jaxtyped, PyTree, TypeCheckError
+
+
+def test_arg_localisation(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: str, y: str, z: int):
+        pass
+
+    matches = [
+        "Type-check error whilst checking the parameters of f",
+        "The problem arose whilst typechecking parameter 'z'.",
+        "Called with parameters: {'x': 'hi', 'y': 'bye', 'z': 'not-an-int'}",
+        r"Parameter annotations: \(x: str, y: str, z: int\).",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f("hi", "bye", "not-an-int")
+
+    @jaxtyped(typechecker=typecheck)
+    def g(x: Float[Array, "a b"], y: Float[Array, "b c"]):
+        pass
+
+    x = jnp.zeros((2, 3))
+    y = jnp.zeros((4, 3))
+    matches = [
+        "Type-check error whilst checking the parameters of g",
+        "The problem arose whilst typechecking parameter 'y'.",
+        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'\]\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        "a=2",
+        "b=3",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            g(x, y=y)
+
+
+def test_return(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    def f(x: PyTree[Any, " T"], y: PyTree[Any, " S"]) -> PyTree[Any, "T S"]:
+        return "foo"
+
+    x = (1, 2)
+    y = {"a": 1}
+    matches = [
+        "Type-check error whilst checking the return value of f",
+        r"Called with parameters: {'x': \(1, 2\), 'y': {'a': 1}}",
+        "Actual value: 'foo'",
+        r"Expected type: PyTree\[Any, \"T S\"\].",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \*\)\)",
+        r"S=PyTreeDef\({'a': \*}\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            f(x, y=y)
+
+
+def test_dataclass_attribute(typecheck):
+    @jaxtyped(typechecker=typecheck)
+    class M(eqx.Module):
+        x: Float[Array, " *foo"]
+        y: PyTree[Any, " T"]
+        z: int
+
+    x = jnp.zeros((2, 3))
+    y = (1, (3, 4))
+    z = "not-an-int"
+
+    matches = [
+        "Type-check error whilst checking the parameters of M",
+        "The problem arose whilst typechecking parameter 'z'.",
+        (
+            r"Called with parameters: {'self': M\(\.\.\.\), 'x': f32\[2,3\], "
+            r"'y': \(1, \(3, 4\)\), 'z': 'not-an-int'}"
+        ),
+        (
+            r"Parameter annotations: \(self: Any, x: Float\[Array, '\*foo'\], "
+            r"y: PyTree\[Any, \"T\"\], z: int\)."
+        ),
+        "The current values for each jaxtyping axis annotation are as follows.",
+        r"foo=\(2, 3\)",
+        (
+            "The current values for each jaxtyping PyTree structure annotation are as "
+            "follows."
+        ),
+        r"T=PyTreeDef\(\(\*, \(\*, \*\)\)\)",
+    ]
+    for match in matches:
+        with pytest.raises(TypeCheckError, match=match):
+            M(x, y, z)

test/test_pytree.py

@@ -25,7 +25,8 @@ import jax.numpy as jnp
 import jax.random as jr
 import pytest
 
-from jaxtyping import Float, jaxtyped, PyTree
+import jaxtyping
+from jaxtyping import AnnotationError, Array, Float, PyTree
 
 from .helpers import make_mlp, ParamError
 
@@ -92,9 +93,8 @@ def test_nested_pytrees(getkey, typecheck):
         g([1, 2, make_mlp()])
 
 
-def test_pytree_array(typecheck):
-    @jaxtyped
-    @typecheck
+def test_pytree_array(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "..."]]):
         pass
 
@@ -106,9 +106,8 @@ def test_pytree_array(typecheck):
         g(1.0)
 
 
-def test_pytree_shaped_array(typecheck, getkey):
-    @jaxtyped
-    @typecheck
+def test_pytree_shaped_array(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
     def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
         pass
 
@@ -193,3 +192,153 @@ def test_subclass_pytree():
     assert issubclass(x, PyTree)
     assert issubclass(y, PyTree)
     assert not issubclass(int, PyTree)
+
+
+def test_structure_match(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, " T"]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+
+    with pytest.raises(ParamError):
+        f(1, ("hi",))
+
+
+def test_structure_prefix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, "T ..."]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+    f(1, ("hi",))
+    f((1, 2), ({"a": "hi"}, {"a": "bye"}))
+    f((1, 2), ({"a": "hi"}, {"not-a": "bye"}))
+
+    with pytest.raises(ParamError):
+        f((1, 2), ({"a": "hi"}, {"a": "bye"}, {"a": "oh-no"}))
+
+    with pytest.raises(ParamError):
+        f((3, 4, 5), {"a": ("hi", "bye")})
+
+
+def test_structure_suffix(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[str, "... T"]):
+        pass
+
+    f(1, "hi")
+    f((3, 4, {"a": 5}), ("a", "b", {"a": "c"}))
+    f(1, ("hi",))
+
+    with pytest.raises(ParamError):
+        f((3, 4), {"a": (1, 2)})
+
+    with pytest.raises(ParamError):
+        f((3, 4, 5), {"a": ("hi", "bye")})
+
+
+def test_structure_compose(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "S T"]):
+        pass
+
+    f(1, 2, "hi")
+    f((1, 2), 2, ("a", "b"))
+
+    with pytest.raises(ParamError):
+        f((1, 2), 2, (1, 2))
+
+    f((1, 2), {"a": 3}, {"a": ("hi", "bye")})
+
+    with pytest.raises(ParamError):
+        f((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
+
+    @jaxtyp(typecheck)
+    def g(x: PyTree[int, " T"], y: PyTree[int, " S"], z: PyTree[str, "T S"]):
+        pass
+
+    with pytest.raises(ParamError):
+        g((1, 2), {"a": 3}, {"a": ("hi", "bye")})
+
+    g((1, 2), {"a": 3}, ({"a": "hi"}, {"a": "bye"}))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_function(variadic, jaxtyp, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyp(typecheck)
+    def f(
+        x: PyTree[Float[Array, jtshape], " T"], y: PyTree[Float[Array, jtshape], " T"]
+    ):
+        pass
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    f(x1, y1)
+    f((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        f(x1, y2)
+
+    with pytest.raises(ParamError):
+        f((x1, x2), (y2, y1))
+
+
+@pytest.mark.parametrize("variadic", (False, True))
+def test_treepath_dependence_dataclass(variadic, typecheck, getkey):
+    if variadic:
+        jtshape = "*?foo"
+        shape = (2, 3)
+    else:
+        jtshape = "?foo"
+        shape = (4,)
+
+    @jaxtyping.jaxtyped(typechecker=typecheck)
+    class A(eqx.Module):
+        x: PyTree[Float[Array, jtshape], " T"]
+        y: PyTree[Float[Array, jtshape], " T"]
+
+    x1 = jr.normal(getkey(), shape)
+    y1 = jr.normal(getkey(), shape)
+    x2 = jr.normal(getkey(), (5,))
+    y2 = jr.normal(getkey(), (5,))
+    A(x1, y1)
+    A((x1, x2), (y1, y2))
+
+    with pytest.raises(ParamError):
+        A(x1, y2)
+
+    with pytest.raises(ParamError):
+        A((x1, x2), (y2, y1))
+
+
+def test_treepath_dependence_missing_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[Float[Array, "?foo"], " T"], y: PyTree[Float[Array, "?foo"]]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    y1 = jr.normal(getkey(), (2,))
+    with pytest.raises(AnnotationError, match="except when contained with structured"):
+        f(x1, y1)
+
+
+def test_treepath_dependence_multiple_structure_annotation(jaxtyp, typecheck, getkey):
+    @jaxtyp(typecheck)
+    def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
+        pass
+
+    x1 = jr.normal(getkey(), (2,))
+    with pytest.raises(AnnotationError, match="ambiguous which PyTree"):
+        f(x1)

test/test_threading.py

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