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

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 ._decorator import jaxtyped as jaxtyped, TypeCheckError as TypeCheckError
+from ._config import config as config
+from ._decorator import jaxtyped as jaxtyped
+from ._errors import (
+    AnnotationError as AnnotationError,
+    TypeCheckError as TypeCheckError,
+)
 from ._import_hook import install_import_hook as install_import_hook
 from ._ipython_extension import load_ipython_extension as load_ipython_extension
 

jaxtyping/_array_types.py

@@ -27,7 +27,7 @@ from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
-from ._raise import jaxtyping_raise, jaxtyping_raise_from
+from ._errors import AnnotationError
 from ._storage import (
     get_shape_memo,
     get_treeflatten_memo,
@@ -133,16 +133,13 @@ def _check_dims(
                 # Make a copy to avoid `__builtins__` getting added as a key.
                 eval_size = eval(elem, single_memo.copy())
             except NameError as e:
-                jaxtyping_raise_from(
+                raise AnnotationError(
-                    NameError(
+                    f"Cannot process symbolic axis '{cls_dim.elem}' as "
-                        f"Cannot process symbolic axis '{cls_dim.elem}' as "
+                    "some axis names have not been processed. In practice you "
-                        "some axis names have not been processed. In practice you "
+                    "should usually only use symbolic axes in annotations "
-                        "should usually only use symbolic axes in annotations "
+                    "for return types, referring only to axes annotated for "
-                        "for return types, referring only to axes annotated for "
+                    "arguments."
-                        "arguments."
+                ) from e
-                    ),
-                    e,
-                )
             if eval_size != obj_size:
                 return False
         else:
@@ -180,8 +177,8 @@ class _MetaAbstractArray(type):
             if len(repr_dtype) == 2 and repr_dtype[0] == "torch":
                 dtype = repr_dtype[1]
             else:
-                jaxtyping_raise(
+                raise AnnotationError(
-                    RuntimeError("Unrecognised array/tensor type to extract dtype from")
+                    "Unrecognised array/tensor type to extract dtype from"
                 )
 
         if cls.dtypes is not _any_dtype:
@@ -561,12 +558,10 @@ def _make_array(array_type, dim_str, dtypes, name):
 
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
-        jaxtyping_raise(
+        raise AnnotationError(
-            RuntimeError(
+            f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
-                f"Do not use `isinstance(x, jaxtyping.{cls.__name__})`. If you want to "
+            "check just the dtype of an array, then use "
-                "check just the dtype of an array, then use "
+            f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
-                f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
-            )
         )
 
     def __getitem__(cls, item: tuple[Any, str]):

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

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

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/_pytree_type.py

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

jaxtyping/_raise.py

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

jaxtyping/_storage.py

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

pyproject.toml

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

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

test/test_array.py

@@ -29,6 +29,7 @@ import torch
 
 from jaxtyping import (
     AbstractDtype,
+    AnnotationError,
     Array,
     ArrayLike,
     Bool,
@@ -448,7 +449,7 @@ def test_incomplete_symbolic(jaxtyp, typecheck, getkey):
         pass
 
     x = jr.normal(getkey(), (4,))
-    with pytest.raises(NameError):
+    with pytest.raises(AnnotationError):
         foo(x)
 
 

test/test_decorator.py

@@ -9,49 +9,49 @@ from .helpers import ParamError, ReturnError
 
 
 class M(metaclass=abc.ABCMeta):
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def f(self):
         ...
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @classmethod
     def g1(cls):
         return 3
 
     @classmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def g2(cls):
         return 4
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @staticmethod
     def h1():
         return 3
 
     @staticmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def h2():
         return 4
 
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @abc.abstractmethod
     def i1(self):
         ...
 
     @abc.abstractmethod
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def i2(self):
         ...
 
 
 class N:
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     @property
     def j1(self):
         return 3
 
     @property
-    @jaxtyped
+    @jaxtyped(typechecker=None)
     def j2(self):
         return 4
 
@@ -154,10 +154,12 @@ def test_local_stringified_annotation(typecheck):
 
     f(LocalFoo())
 
-    @jaxtyped
+    with pytest.warns(match="As of jaxtyping version 0.2.24"):
-    @typecheck
+
-    def g(x: "LocalFoo") -> "LocalFoo":
+        @jaxtyped
-        return x
+        @typecheck
+        def g(x: "LocalFoo") -> "LocalFoo":
+            return x
 
     g(LocalFoo())
 

test/test_messages.py

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

test/test_pytree.py

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

test/test_threading.py

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