Some improvements (#77)

* Various improvements.

- Added support for functions in symbolic dimensions, e.g. "min(foo,bar)", which were previously disallowed due to the presence of a comma. (#51)
- Added support for adding ignored names to dimensions, e.g. "cols=4". (#76)

* Now works with Python 3.10 A | B union types.

.flake8

@@ -1,4 +0,0 @@
-[flake8]
-max-line-length = 88
-ignore = W291,W293,W503,W504,E123,E126,E203,E402,E701,E731,F722
-per-file-ignores = __init__.py: F401

.github/workflows/release.yml

@@ -33,7 +33,8 @@ jobs:
         with:
             python-version: "3.8"
             test-script: |
-                python -m pip install pytest beartype equinox jaxlib
+                python -m pip install pytest beartype equinox jaxlib cloudpickle
+                python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
                 cp -r ${{ github.workspace }}/test ./test
                 pytest
             pypi-token: ${{ secrets.pypi_token }}

.github/workflows/run_tests.yml

@@ -26,7 +26,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ 3.8, 3.9 ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
@@ -42,7 +42,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest wheel beartype equinox jaxlib
+          python -m pip install pytest wheel beartype equinox jaxlib cloudpickle
+          python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
 
       - name: Checks with pre-commit
         uses: pre-commit/action@v2.0.3

.isort.cfg

@@ -1,6 +0,0 @@
-[settings]
-force_alphabetical_sort_within_sections=true
-lines_after_imports=2
-profile=black
-treat_comments_as_code=true
-extra_standard_library=typing_extensions

.pre-commit-config.yaml

@@ -22,17 +22,7 @@ repos:
     rev: 22.3.0
     hooks:
     - id: black
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.2.3
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: 'v0.0.255'
     hooks:
-    - id: nbqa-black
-    - id: nbqa-isort 
-    - id: nbqa-flake8
-  - repo: https://github.com/PyCQA/isort
-    rev: 5.10.1
-    hooks:
-    - id: isort
-  - repo: https://github.com/pycqa/flake8
-    rev: 4.0.1
-    hooks:
-    - id: flake8
+      - id: ruff

API.md

@@ -17,6 +17,7 @@ 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. `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 _ _"`.
+- 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"]`.
 
 When using multiple modifiers, their order does not matter.
 
@@ -62,6 +63,8 @@ Float32[Array, "some_shape"]
 
 The array should typically be a `jaxtyping.Array`, which is an alias for `jax.numpy.ndarray`.
 
+`jaxtyping.ArrayLike` is also available, which is an alias for `jax.typing.ArrayLike`. This is a union over JAX arrays and the builtin `bool`/`int`/`float`/`complex`.
+
 But you can use other types as well. `jaxtyping` has support for JAX, NumPy, TensorFlow, and PyTorch, e.g.:
 ```python
 Float[np.ndarray, "..."]
@@ -142,13 +145,13 @@ from jaxtyping import install_import_hook
 # Plus any one of the following:
 
 # decorate @jaxtyped and @typeguard.typechecked
-with install_import_hook("foo", ("typeguard", "typechecked")):
+with install_import_hook("foo", "typeguard.typechecked"):
     import foo          # Any module imported inside this `with` block, whose name begins
     import foo.bar      # with the specified string, will automatically have both `@jaxtyped`
     import foo.bar.qux  # and the specified typechecker applied to all of their functions.
 
 # decorate @jaxtyped and @beartype.beartype
-with install_import_hook("foo", ("beartype", "beartype")):
+with install_import_hook("foo", "beartype.beartype"):
     ...
     
 # decorate only @jaxtyped (if you want that for some reason)
@@ -168,12 +171,14 @@ The import hook can be applied to multiple packages via
 install_import_hook(["foo", "bar.baz"], ...)
 ```
 
+The import hook will automatically decorate all functions, and the `__init__` method of dataclasses.
+
 **Example: writing an end-user script**
 
 ```python
 ### entry_point.py
 from jaxtyping import install_import_hook
-with install_import_hook("do_stuff", ("typeguard", "typechecked")):
+with install_import_hook("do_stuff", "typeguard.typechecked"):
     import do_stuff
 
 ### do_stuff.py
@@ -188,7 +193,7 @@ def g(x: Float32[Array, "..."]):
 ```python
 ### __init__.py
 from jaxtyping import install_import_hook
-with install_import_hook("my_library_name", ("beartype", "beartype")):
+with install_import_hook("my_library_name", "beartype.beartype"):
     from .subpackage import foo  # full name is my_library_name.subpackage so will be hook'd
     from .another_subpackage import bar  # full name is my_library_name.another_subpackage so will be hook'd.
 ```

CONTRIBUTING.md

@@ -28,7 +28,8 @@ Now make your changes. Make sure to include additional tests if necessary.
 Next verify the tests all pass:
 
 ```bash
-pip install pytest
+pip install pytest cloudpickle
+pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
 pytest
 ```
 

FAQ.md

@@ -12,13 +12,13 @@ jaxtyping and `jax.jit` synergise beautifully.
 
 When calling JAX operations wrapped in a `jax.jit`, then the dtype/shape-checking will happen at trace time. (When JAX traces your function prior to compiling it.) The actual compiled code does not have any dtype/shape-checking, and will therefore still be just as fast as before!
 
-## `flake8` is throwing an error.
+## `flake8` or Ruff are throwing an error.
 
 In type annotations, strings are used for two different things. Sometimes they're strings. Sometimes they're "forward references", used to refer to a type that will be defined later.
 
 Some tooling in the Python ecosystem assumes that only the latter is true, and will throw spurious errors if you try to use a string just as a string (like we do).
 
-In the case of `flake8`, at least, this is easily resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
+In the case of `flake8`, or Ruff, this can be resolved. Multi-dimensional arrays (e.g. `Float32[Array, "b c"]`) will throw a very unusual error (F722, syntax error in forward annotation), so you can safely just disable this particular error globally. Uni-dimensional arrays (e.g. `Float32[Array, "x"]`) will throw an error that's actually useful (F821, undefined name), so instead of disabling this globally, you should instead prepend a space to the start of your shape, e.g. `Float32[Array, " x"]`. `jaxtyping` will treat this in the same way, whilst `flake8` will now throw an F722 error that you can disable as before.
 
 ## Does jaxtyping use [PEP 646](https://www.python.org/dev/peps/pep-0646/) (variadic generics)?
 

MANIFEST.in

@@ -1,2 +1,2 @@
-include LICENSE
-prune tests
+include jaxtyping/py.typed
+prune test

README.md

@@ -2,9 +2,10 @@
 
 Type annotations **and runtime checking** for:
 
-1. shape and dtype of [JAX](https://github.com/google/jax) arrays;
+1. shape and dtype of [JAX](https://github.com/google/jax) arrays; *(Now also supports PyTorch, NumPy, and TensorFlow!)*
 2. [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html).
 
+
 **For example:**
 ```python
 from jaxtyping import Array, Float, PyTree
@@ -28,7 +29,9 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires JAX 0.3.4+.
+Requires Python 3.8+.
+
+JAX is an optional dependency, required for `jaxtyping.{Array, ArrayLike, PyTree}`. If JAX is not installed then these types will not be available, but you may still use jaxtyping alongside PyTorch/NumPy/etc.
 
 Also install your favourite runtime type-checking package. The two most popular are [typeguard](https://github.com/agronholm/typeguard) (which exhaustively checks every argument) and [beartype](https://github.com/beartype/beartype) (which checks random pieces of arguments).
 
@@ -46,14 +49,10 @@ Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
 
 Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
 
+Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+
 SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
 
-### Acknowledgements
-
-Shape annotations + runtime type checking is inspired by [TorchTyping](https://github.com/patrick-kidger/torchtyping).
-
-The concise syntax is partially inspired by [etils.array_types](https://github.com/google/etils/tree/main/etils/array_types).
-
 ### Disclaimer
 
 This is not an official Google product.

jaxtyping/__init__.py

@@ -17,40 +17,101 @@
 # 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.
 
-from jax.numpy import ndarray as Array
+import typing
+
+
+try:
+    import jax
+except ImportError:
+    has_jax = False
+else:
+    has_jax = True
+    del jax
+
+
+# Type checkers don't know which branch below will be executed.
+if typing.TYPE_CHECKING:
+    # For imports, we need to explicitly `import X as X` in order for Pyright to see
+    # them as public. See discussion at https://github.com/microsoft/pyright/issues/2277
+    from jax import Array as Array
+    from jax.typing import ArrayLike as ArrayLike
+elif has_jax:
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+        class Array:
+            pass
+
+        Array.__module__ = "builtins"
+
+        class ArrayLike:
+            pass
+
+        ArrayLike.__module__ = "builtins"
+    else:
+        from jax import Array as Array
+
+        try:
+            from jax.typing import ArrayLike as ArrayLike
+        except (ModuleNotFoundError, ImportError):
+            pass
 
 from .array_types import (
-    AbstractArray,
-    AbstractDtype,
-    BFloat16,
-    Bool,
-    Complex,
-    Complex64,
-    Complex128,
-    Float,
-    Float16,
-    Float32,
-    Float64,
-    get_array_name_format,
-    Inexact,
-    Int,
-    Int8,
-    Int16,
-    Int32,
-    Int64,
-    Integer,
-    Num,
-    set_array_name_format,
-    Shaped,
-    UInt,
-    UInt8,
-    UInt16,
-    UInt32,
-    UInt64,
+    AbstractArray as AbstractArray,
+    AbstractDtype as AbstractDtype,
+    BFloat16 as BFloat16,
+    Bool as Bool,
+    Complex as Complex,
+    Complex64 as Complex64,
+    Complex128 as Complex128,
+    Float as Float,
+    Float16 as Float16,
+    Float32 as Float32,
+    Float64 as Float64,
+    get_array_name_format as get_array_name_format,
+    Inexact as Inexact,
+    Int as Int,
+    Int8 as Int8,
+    Int16 as Int16,
+    Int32 as Int32,
+    Int64 as Int64,
+    Integer as Integer,
+    Num as Num,
+    set_array_name_format as set_array_name_format,
+    Shaped as Shaped,
+    UInt as UInt,
+    UInt8 as UInt8,
+    UInt16 as UInt16,
+    UInt32 as UInt32,
+    UInt64 as UInt64,
 )
-from .decorator import jaxtyped
-from .import_hook import install_import_hook
-from .pytree_type import PyTree
+from .decorator import jaxtyped as jaxtyped
+from .import_hook import install_import_hook as install_import_hook
 
 
-__version__ = "0.2.1"
+if typing.TYPE_CHECKING:
+    # Set up to deliberately confuse a static type checker.
+    import typing_extensions
+
+    PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar")
+    # What's going on with this madness?
+    #
+    # At static-type-checking-time, we want `PyTree` to be a type for which both
+    # `PyTree` and `PyTree[Foo]` are equivalent to `Any`.
+    # (The intention is that `PyTree` be a runtime-only type; there's no real way to
+    # do more with static type checkers.)
+    #
+    # Unfortunately, this isn't possible: `Any` isn't subscriptable. And there's no
+    # equivalent way we can fake this using typing annotations. (In some sense the
+    # closest thing would be a `Protocol[T]` with no methods, but that's actually the
+    # opposite of what we want: that ends up allowing nothing at all.)
+    #
+    # The good news for us is that static type checkers have an internal escape hatch.
+    # If they can't figure out what a type is, then they just give up and allow
+    # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
+    # annotation, which static type checkers aren't smart enough to resolve.
+elif has_jax:
+    from .pytree_type import PyTree as PyTree  # noqa: F401
+
+del has_jax
+
+__version__ = "0.2.15"

jaxtyping/array_types.py

@@ -19,8 +19,20 @@
 
 import enum
 import functools as ft
-from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union
-from typing_extensions import Literal
+import sys
+import types
+import typing
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    NoReturn,
+    Optional,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import numpy as np
 
@@ -149,25 +161,25 @@ class _MetaAbstractArray(type):
         if cls.dtypes is not _any_dtype and dtype not in cls.dtypes:
             return False
 
-        if len(storage.memo_stack) == 0:
+        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 = {}
-            temp_memo = True
-        else:
-            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()
-            temp_memo = False
 
         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 not temp_memo:
+            if no_temp_memo:
                 storage.memo_stack[-1] = (
                     single_memo,
                     variadic_memo,
@@ -233,6 +245,20 @@ class _MetaAbstractArray(type):
             assert False
 
 
+@ft.lru_cache(maxsize=None)
+def _make_metaclass(base_metaclass):
+    class MetaAbstractArray(_MetaAbstractArray, base_metaclass):
+        pass
+
+    return MetaAbstractArray
+
+
+def _check_scalar(dtype, dtypes, dims):
+    if len(dims) != 0:
+        return dims == (_anonymous_variadic_dim,)
+    return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
+
+
 class AbstractArray(metaclass=_MetaAbstractArray):
     array_type: Any
     dtypes: List[str]
@@ -240,6 +266,197 @@ class AbstractArray(metaclass=_MetaAbstractArray):
     index_variadic: Optional[int]
 
 
+_not_made = object()
+
+
+_union_types = [typing.Union]
+if sys.version_info >= (3, 10):
+    _union_types.append(types.UnionType)
+
+
+@ft.lru_cache(maxsize=None)
+def _make_array(array_type, dim_str, dtypes, name):
+    if not isinstance(dim_str, str):
+        raise ValueError(
+            "Shape specification must be a string. Axes should be separated with "
+            "spaces."
+        )
+    dims = []
+    index_variadic = None
+    for index, elem in enumerate(dim_str.split()):
+        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")
+        if elem.endswith("#"):
+            raise ValueError(
+                "As of jaxtyping v0.1.0, broadcastable dimensions 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; "
+                    f"got {elem}"
+                )
+            broadcastable = False
+            variadic = True
+            anonymous = True
+            dim_type = _DimType.named
+        else:
+            broadcastable = False
+            variadic = False
+            anonymous = False
+            while True:
+                if len(elem) == 0:
+                    # This branch needed as just `_` is valid
+                    break
+                first_char = elem[0]
+                if first_char == "#":
+                    if broadcastable:
+                        raise ValueError(
+                            "Do not use # twice to denote broadcastability, e.g. "
+                            "`##foo` is not allowed"
+                        )
+                    broadcastable = True
+                    elem = elem[1:]
+                elif first_char == "*":
+                    if variadic:
+                        raise ValueError(
+                            "Do not use * twice to denote accepting multiple "
+                            "dimensions, e.g. `**foo` is not allowed"
+                        )
+                    variadic = True
+                    elem = elem[1:]
+                elif first_char == "_":
+                    if anonymous:
+                        raise ValueError(
+                            "Do not use _ twice to denote anonymity, e.g. `__foo` "
+                            "is not allowed"
+                        )
+                    anonymous = 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:
+                    _, elem = elem.split("=")
+                else:
+                    break
+            if len(elem) == 0 or elem.isidentifier():
+                dim_type = _DimType.named
+            else:
+                try:
+                    elem = int(elem)
+                except ValueError:
+                    dim_type = _DimType.symbolic
+                else:
+                    dim_type = _DimType.fixed
+
+        if variadic:
+            if index_variadic is not None:
+                raise ValueError(
+                    "Cannot use multiple-dimension 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"
+                )
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a fixed axis be anonymous, 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"
+                    )
+                if variadic:
+                    elem = _anonymous_variadic_dim
+                else:
+                    elem = _anonymous_dim
+            else:
+                if variadic:
+                    elem = _NamedVariadicDim(elem, broadcastable)
+                else:
+                    elem = _NamedDim(elem, broadcastable)
+        else:
+            assert dim_type is _DimType.symbolic
+            if anonymous:
+                raise ValueError(
+                    "Cannot have a symbolic dimension be anonymous, e.g. "
+                    "`_foo+bar` is not allowed"
+                )
+            if variadic:
+                raise ValueError(
+                    "Cannot have symbolic multiple-dimensions, e.g. "
+                    "`*foo+bar` is not allowed"
+                )
+            elem = compile(elem, "<string>", "eval")
+            elem = _SymbolicDim(elem, broadcastable)
+        dims.append(elem)
+    dims = tuple(dims)
+
+    # Allow Python built-in numeric types.
+    # TODO: do something more generic than this? Should we _make all types
+    # that have `shape` and `dtype` attributes or something?
+    if array_type is bool:
+        if _check_scalar("bool", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is int:
+        if _check_scalar("int", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is float:
+        if _check_scalar("float", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    elif array_type is complex:
+        if _check_scalar("complex", dtypes, dims):
+            return array_type
+        else:
+            return _not_made
+    try:
+        type_str = array_type.__name__
+    except AttributeError:
+        type_str = repr(array_type)
+    if _array_name_format == "dtype_and_shape":
+        name = f"{name}[{type_str}, '{dim_str}']"
+    elif _array_name_format == "array":
+        name = type_str
+    else:
+        raise ValueError(f"array_name_format {_array_name_format} not recognised")
+    metaclass = _make_metaclass(type(array_type))
+    out = metaclass(
+        name,
+        (array_type, AbstractArray),
+        dict(
+            array_type=array_type,
+            dtypes=dtypes,
+            dims=dims,
+            index_variadic=index_variadic,
+        ),
+    )
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        out.__module__ = "builtins"
+    else:
+        out.__module__ = "jaxtyping"
+    return out
+
+
 class _MetaAbstractDtype(type):
     def __instancecheck__(cls, obj: Any) -> NoReturn:
         raise RuntimeError(
@@ -248,8 +465,7 @@ class _MetaAbstractDtype(type):
             f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
         )
 
-    @ft.lru_cache(maxsize=None)
-    def __getitem__(cls, item: Tuple[Any, str]) -> _MetaAbstractArray:
+    def __getitem__(cls, item: Tuple[Any, str]):
         if not isinstance(item, tuple) or len(item) != 2:
             raise ValueError(
                 "As of jaxtyping v0.2.0, type annotations must now include an explicit "
@@ -257,147 +473,18 @@ class _MetaAbstractDtype(type):
             )
         array_type, dim_str = item
         del item
-        if not isinstance(dim_str, str):
-            raise ValueError(
-                "Shape specification must be a string. Axes should be separated with "
-                "spaces."
-            )
-        dims = []
-        index_variadic = None
-        for index, elem in enumerate(dim_str.split()):
-            if "," in elem:
-                # Common mistake
-                raise ValueError(
-                    "Dimensions should be separated with spaces, not commas"
-                )
-            if elem.endswith("#"):
-                raise ValueError(
-                    "As of jaxtyping v0.1.0, broadcastable dimensions 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; "
-                        f"got {elem}"
-                    )
-                broadcastable = False
-                variadic = True
-                anonymous = True
-                dim_type = _DimType.named
-            else:
-                broadcastable = False
-                variadic = False
-                anonymous = False
-                while True:
-                    if len(elem) == 0:
-                        # This branch needed as just `_` is valid
-                        break
-                    first_char = elem[0]
-                    if first_char == "#":
-                        if broadcastable:
-                            raise ValueError(
-                                "Do not use # twice to denote broadcastability, e.g. "
-                                "`##foo` is not allowed"
-                            )
-                        broadcastable = True
-                        elem = elem[1:]
-                    elif first_char == "*":
-                        if variadic:
-                            raise ValueError(
-                                "Do not use * twice to denote accepting multiple "
-                                "dimensions, e.g. `**foo` is not allowed"
-                            )
-                        variadic = True
-                        elem = elem[1:]
-                    elif first_char == "_":
-                        if anonymous:
-                            raise ValueError(
-                                "Do not use _ twice to denote anonymity, e.g. `__foo` "
-                                "is not allowed"
-                            )
-                        anonymous = True
-                        elem = elem[1:]
-                    else:
-                        break
-                try:
-                    elem = int(elem)
-                except ValueError:
-                    if len(elem) == 0 or elem.isidentifier():
-                        dim_type = _DimType.named
-                    else:
-                        dim_type = _DimType.symbolic
-                else:
-                    dim_type = _DimType.fixed
-
-            if variadic:
-                if index_variadic is not None:
-                    raise ValueError(
-                        "Cannot use multiple-dimension 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"
-                    )
-                if anonymous:
-                    raise ValueError(
-                        "Cannot have a fixed axis be anonymous, 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"
-                        )
-                    if variadic:
-                        elem = _anonymous_variadic_dim
-                    else:
-                        elem = _anonymous_dim
-                else:
-                    if variadic:
-                        elem = _NamedVariadicDim(elem, broadcastable)
-                    else:
-                        elem = _NamedDim(elem, broadcastable)
-            else:
-                assert dim_type is _DimType.symbolic
-                if anonymous:
-                    raise ValueError(
-                        "Cannot have a symbolic dimension be anonymous, e.g. "
-                        "`_foo+bar` is not allowed"
-                    )
-                if variadic:
-                    raise ValueError(
-                        "Cannot have symbolic multiple-dimensions, e.g. "
-                        "`*foo+bar` is not allowed"
-                    )
-                elem = compile(elem, "<string>", "eval")
-                elem = _SymbolicDim(elem, broadcastable)
-            dims.append(elem)
-        if _array_name_format == "dtype_and_shape":
-            name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
-        elif _array_name_format == "array":
-            name = "Array"
+        if typing.get_origin(array_type) in _union_types:
+            out = [
+                _make_array(x, dim_str, cls.dtypes, cls.__name__)
+                for x in typing.get_args(array_type)
+            ]
+            out = tuple(x for x in out if x is not _not_made)
+            out = Union[out]
         else:
-            raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        return _MetaAbstractArray(
-            name,
-            (AbstractArray,),
-            dict(
-                array_type=array_type,
-                dtypes=cls.dtypes,
-                dims=dims,
-                index_variadic=index_variadic,
-            ),
-        )
+            out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__)
+            if out is _not_made:
+                raise ValueError("Invalid jaxtyping type annotation.")
+        return out
 
 
 class AbstractDtype(metaclass=_MetaAbstractDtype):
@@ -414,7 +501,9 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
 
         dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes
         if isinstance(dtypes, str):
-            dtypes = [dtypes]
+            dtypes = (dtypes,)
+        elif dtypes is not _any_dtype:
+            dtypes = tuple(dtypes)
         cls.dtypes = dtypes
 
 
@@ -422,31 +511,34 @@ if TYPE_CHECKING:
     # Note that `from typing_extensions import Annotated; ... = Annotated`
     # does not work with static type checkers. `Annotated` is a typeform rather
     # than a type, meaning it cannot be assigned.
-    from typing_extensions import Annotated as BFloat16
-    from typing_extensions import Annotated as Bool
-    from typing_extensions import Annotated as Complex
-    from typing_extensions import Annotated as Complex64
-    from typing_extensions import Annotated as Complex128
-    from typing_extensions import Annotated as Float
-    from typing_extensions import Annotated as Float16
-    from typing_extensions import Annotated as Float32
-    from typing_extensions import Annotated as Float64
-    from typing_extensions import Annotated as Inexact
-    from typing_extensions import Annotated as Int
-    from typing_extensions import Annotated as Int8
-    from typing_extensions import Annotated as Int16
-    from typing_extensions import Annotated as Int32
-    from typing_extensions import Annotated as Int64
-    from typing_extensions import Annotated as Integer
-    from typing_extensions import Annotated as Num
-    from typing_extensions import Annotated as Shaped
-    from typing_extensions import Annotated as UInt
-    from typing_extensions import Annotated as UInt8
-    from typing_extensions import Annotated as UInt16
-    from typing_extensions import Annotated as UInt32
-    from typing_extensions import Annotated as UInt64
+    from typing_extensions import (
+        Annotated as BFloat16,
+        Annotated as Bool,
+        Annotated as Complex,
+        Annotated as Complex64,
+        Annotated as Complex128,
+        Annotated as Float,
+        Annotated as Float16,
+        Annotated as Float32,
+        Annotated as Float64,
+        Annotated as Inexact,
+        Annotated as Int,
+        Annotated as Int8,
+        Annotated as Int16,
+        Annotated as Int32,
+        Annotated as Int64,
+        Annotated as Integer,
+        Annotated as Num,
+        Annotated as Shaped,
+        Annotated as UInt,
+        Annotated as UInt8,
+        Annotated as UInt16,
+        Annotated as UInt32,
+        Annotated as UInt64,
+    )
 else:
-    _bool = "bool_"
+    _bool = "bool"
+    _bool_ = "bool_"
     _uint8 = "uint8"
     _uint16 = "uint16"
     _uint32 = "uint32"
@@ -468,6 +560,10 @@ else:
 
         _Cls.__name__ = name
         _Cls.__qualname__ = name
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            _Cls.__module__ = "builtins"
+        else:
+            _Cls.__module__ = "jaxtyping"
         return _Cls
 
     UInt8 = _make_dtype(_uint8, "UInt8")
@@ -485,6 +581,7 @@ else:
     Complex64 = _make_dtype(_complex64, "Complex64")
     Complex128 = _make_dtype(_complex128, "Complex128")
 
+    bools = [_bool, _bool_]
     uints = [_uint8, _uint16, _uint32, _uint64]
     ints = [_int8, _int16, _int32, _int64]
     floats = [_bfloat16, _float16, _float32, _float64]
@@ -493,7 +590,7 @@ else:
     # We match NumPy's type hierarachy in what types to provide. See the diagram at
     # https://numpy.org/doc/stable/reference/arrays.scalars.html#scalars
 
-    Bool = _make_dtype(_bool, "Bool")
+    Bool = _make_dtype(bools, "Bool")
     UInt = _make_dtype(uints, "UInt")
     Int = _make_dtype(ints, "Int")
     Integer = _make_dtype(uints + ints, "Integer")

jaxtyping/decorator.py

@@ -17,21 +17,84 @@
 # 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
 import functools as ft
+import inspect
 import threading
+import types
+import weakref
 
 
 storage = threading.local()
-storage.memo_stack = []
+
+
+_jaxtyped_fns = weakref.WeakSet()
 
 
 def jaxtyped(fn):
-    @ft.wraps(fn)
-    def wrapper(*args, **kwargs):
-        storage.memo_stack.append(({}, {}, {}))
-        try:
-            return fn(*args, **kwargs)
-        finally:
-            storage.memo_stack.pop()
+    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:
+            raise ValueError(
+                "jaxtyped may only be added as a class decorator to dataclasses"
+            )
+    # 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__))
+    elif isinstance(fn, staticmethod):
+        return staticmethod(jaxtyped(fn.__func__))
+    elif isinstance(fn, property):
+        if fn.fget is None:
+            fget = None
+        else:
+            fget = jaxtyped(fn.fget)
+        if fn.fset is None:
+            fset = None
+        else:
+            fset = jaxtyped(fn.fset)
+        if fn.fdel is None:
+            fdel = None
+        else:
+            fdel = jaxtyped(fn.fdel)
+        return property(fget=fget, fset=fset, fdel=fdel)
+    else:
 
-    return wrapper
+        @ft.wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            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()
+
+        _jaxtyped_fns.add(wrapped_fn)
+        return wrapped_fn
+
+
+def _jaxtyped_typechecker(typechecker):
+    # typechecker is expected to probably be either `typeguard.typechecked`, or
+    # `beartype.beartype`, or `None`.
+
+    if typechecker is None:
+        typechecker = lambda x: x
+
+    def _wrapper(kls):
+        assert inspect.isclass(kls)
+        if dataclasses.is_dataclass(kls):
+            init = jaxtyped(typechecker(kls.__init__))
+            kls.__init__ = init
+        return kls
+
+    return _wrapper

jaxtyping/import_hook.py

@@ -50,12 +50,13 @@
 
 
 import ast
+import functools as ft
 import sys
 from importlib.abc import MetaPathFinder
 from importlib.machinery import SourceFileLoader
 from importlib.util import cache_from_source, decode_source
 from inspect import isclass
-from typing import Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple, Union
 from unittest.mock import patch
 
 
@@ -64,8 +65,31 @@ def _call_with_frames_removed(f, *args, **kwargs):
     return f(*args, **kwargs)
 
 
-def _optimized_cache_from_source(path, debug_override=None):
-    return cache_from_source(path, debug_override, optimization="jaxtyping")
+def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None):
+    # Version 2: change the position of the `@jaxtyped` decorator, so need a
+    #     different name to avoid hitting old __pycache__.
+    # Version 3: now also annotating classes.
+    # Version 4: I'm honestly not sure, but bumping this fixed some kind of odd error.
+    #     Maybe I changed something with hte classes part way through version 3?
+    # Version 5: Added support for string-based `typechecker` argument.
+    # Version 6: optimization tag now depends on `typechecker` argument, so that
+    #    changing the typechecker will hit a different cache.
+    return cache_from_source(
+        path, debug_override, optimization=f"jaxtyping6{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
+
+
+def _str_lookup(string):
+    module = ast.parse(string)
+    (expr,) = module.body
+    return expr.value
 
 
 class _JaxtypingTransformer(ast.NodeVisitor):
@@ -84,7 +108,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
             else:
                 node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)]))
                 if self._typechecker is not None:
-                    typechecker_module, _ = self._typechecker
+                    typechecker_module, _ = self._typechecker.split(".", 1)
                     node.body.insert(
                         i, ast.Import(names=[ast.alias(typechecker_module, None)])
                     )
@@ -95,31 +119,37 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         self._parents.pop()
         return node
 
+    def visit_ClassDef(self, node: ast.ClassDef):
+        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
+        if self._typechecker is None:
+            args = [ast.Constant(None)]
+        else:
+            args = [_str_lookup(self._typechecker)]
+        node.decorator_list.insert(0, ast.Call(func, args, keywords=[]))
+        self._parents.append(node)
+        self.generic_visit(node)
+        self._parents.pop()
+        return node
+
     def visit_FunctionDef(self, node: ast.FunctionDef):
         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 start of the decorator list, in case a typechecking
-            # annotation has been manually applied; we need to be above that.
-            node.decorator_list.insert(
-                0,
-                ast.Attribute(
-                    ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load()
-                ),
-            )
+            # 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.
+            #
+            # 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
+            # has been manually applied, and we'd need to be above that. In this
+            # 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.
+            node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
             if self._typechecker is not None:
                 # Place at the end of the decorator list, as decorators
                 # frequently remove annotations from functions and we'd like to
                 # use those annotations.
-                typechecker_module, typechecker_function = self._typechecker
-                node.decorator_list.append(
-                    ast.Attribute(
-                        ast.Name(id=typechecker_module, ctx=ast.Load()),
-                        typechecker_function,
-                        ast.Load(),
-                    )
-                )
-
+                node.decorator_list.append(_str_lookup(self._typechecker))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -130,6 +160,7 @@ class _JaxtypingLoader(SourceFileLoader):
     def __init__(self, *args, typechecker, **kwargs):
         super().__init__(*args, **kwargs)
         self._typechecker = typechecker
+        self._typechecker_hash = str(abs(hash(self._typechecker)))
 
     def source_to_code(self, data, path, *, _optimize=-1):
         source = decode_source(data)
@@ -153,7 +184,7 @@ class _JaxtypingLoader(SourceFileLoader):
         # patch safe
         with patch(
             "importlib._bootstrap_external.cache_from_source",
-            _optimized_cache_from_source,
+            ft.partial(_optimized_cache_from_source, self._typechecker_hash),
         ):
             return super().exec_module(module)
 
@@ -216,7 +247,7 @@ class ImportHookManager:
 # Deliberately no default for `typechecker` so that folks must opt-in to not having
 # a typechecker.
 def install_import_hook(
-    modules: Iterable[str], typechecker: Optional[Tuple[str, str]]
+    modules: Iterable[str], typechecker: Optional[Union[str, Tuple[str, str]]]
 ) -> ImportHookManager:
     """Automatically apply `@jaxtyped`, and optionally a type checker, to all classes
     and functions.
@@ -228,10 +259,18 @@ def install_import_hook(
     - `packages`: the names of the modules in which to automatically apply `@jaxtyped`
         and `@typechecked`.
     - `typechecker`: the module and function of the typechecker you want to use, as a
-        2-tuple of strings. For example `typechecker=("typeguard", "typechecked")` or
-        `typechecker=("beartype", "beartype")`. You may pass `typechecker=None` if you
-        do not want to automatically decorate with a typechecker as well; e.g. if you
-        have a codebase that already has these decorators.
+        string. For example `typechecker="typeguard.typechecked"`, or
+        `typechecker="beartype.beartype"`. You may pass `typechecker=None` if you do not
+        want to automatically decorate with a typechecker as well.
+
+    If the function already has any decorators on it, then both the `@jaxtyped` and the
+    typechecker decorators will go at the bottom of the decorator list, e.g.
+    ```python
+    @some_other_decorator
+    @jaxtyped
+    @beartype.beartype
+    def foo(...): ...
+    ```
 
     **Returns:**
 
@@ -243,8 +282,8 @@ def install_import_hook(
     ```python
     # entry_point.py
     from jaxtyped import install_import_hook
-    install_import_hook("main", ("beartype", "beartype"))
-    import main
+    with install_import_hook("main", ("beartype", "beartype"))
+        import main
     ...  # do whatever you're doing
 
     # main.py
@@ -260,6 +299,10 @@ def install_import_hook(
     if isinstance(modules, str):
         modules = [modules]
 
+    # Support old less-flexible API.
+    if isinstance(typechecker, tuple):
+        typechecker = ".".join(typechecker)
+
     for i, finder in enumerate(sys.meta_path):
         if (
             isclass(finder)

jaxtyping/pytest_plugin.py

@@ -52,4 +52,4 @@ def pytest_configure(config):
         )
         raise RuntimeError(message.format(", ".join(already_imported_packages)))
 
-    install_import_hook(packages, typechecker.rsplit(".", 1))
+    install_import_hook(packages, typechecker)

jaxtyping/pytree_type.py

@@ -18,9 +18,10 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+import typing
 from typing import Generic, TypeVar
 
-import jax
+import jax.tree_util as jtu
 import typeguard
 
 
@@ -33,6 +34,7 @@ class _FakePyTree(Generic[_T]):
 
 _FakePyTree.__name__ = "PyTree"
 _FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
 # Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms
 # isn't allowed.
 # Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__
@@ -49,7 +51,20 @@ class _MetaPyTree(type):
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
         name = str(_FakePyTree[item])
-        return _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
+        return out
+
+
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
 
 
 class _MetaSubscriptPyTree(type):
@@ -70,16 +85,20 @@ class _MetaSubscriptPyTree(type):
         def is_leaftype(x):
             try:
                 accepts_leaftype(x)
-            except TypeError:
+            except _TypeCheckError:
                 return False
             else:
                 return True
 
-        leaves = jax.tree_leaves(obj, is_leaf=is_leaftype)
+        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
         return all(map(is_leaftype, leaves))
 
 
-PyTree = _MetaPyTree("PyTree", (), {})
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
-# instancecheck for PyTree[foo], but we subclassing
+# instancecheck for PyTree[foo], but subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.
+PyTree = _MetaPyTree("PyTree", (), {})
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    PyTree.__module__ = "builtins"
+else:
+    PyTree.__module__ = "jaxtyping"

pyproject.toml

@@ -0,0 +1,10 @@
+[tool.ruff]
+select = ["E", "F", "I001"]
+ignore = ["E721", "E731", "F722"]
+ignore-init-module-imports = true
+
+[tool.ruff.isort]
+combine-as-imports = true
+lines-after-imports = 2
+extra-standard-library = ["typing_extensions"]
+order-by-type = false

setup.py

@@ -63,15 +63,17 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 
-python_requires = "~=3.7"
+python_requires = "~=3.8"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.
+# For typing_extensions, we choose versions that match
+# https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab
+
 install_requires = [
-    "jax>=0.3.4",
     "numpy>=1.20.0",
     "typeguard>=2.13.3",
-    "typing_extensions>=4.2.0",
+    "typing_extensions>=3.7.4.1",
 ]
 
 entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])
@@ -94,4 +96,5 @@ setuptools.setup(
     install_requires=install_requires,
     entry_points=entry_points,
     packages=[name],
+    include_package_data=True,
 )

test/helpers.py

@@ -18,16 +18,31 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import equinox as eqx
+import typeguard
 
 
+ParamError = []
+ReturnError = []
+ParamError.append(TypeError)  # old typeguard
+ReturnError.append(TypeError)  # old typeguard
+
+try:
+    # new typeguard
+    ParamError.append(typeguard.TypeCheckError)
+    ReturnError.append(typeguard.TypeCheckError)
+except AttributeError:
+    pass
+
 try:
     import beartype
 except ImportError:
-    ParamError = TypeError
-    ReturnError = TypeError
+    pass
 else:
-    ParamError = (TypeError, beartype.roar.BeartypeCallHintParamViolation)
-    ReturnError = (TypeError, beartype.roar.BeartypeCallHintReturnViolation)
+    ParamError.append(beartype.roar.BeartypeCallHintParamViolation)
+    ReturnError.append(beartype.roar.BeartypeCallHintReturnViolation)
+
+ParamError = tuple(ParamError)
+ReturnError = tuple(ReturnError)
 
 
 @eqx.filter_jit

test/import_hook_tester_beartype.py

@@ -17,6 +17,9 @@
 # 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
+
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +35,28 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_beartype_full.py

@@ -0,0 +1,62 @@
+# 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 dataclasses
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_beartype_old.py

@@ -0,0 +1,62 @@
+# 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 dataclasses
+
+import equinox as eqx
+import jax.numpy as jnp
+import pytest
+
+from jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_broken_checker.py

@@ -17,6 +17,9 @@
 # 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
+
+import equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +35,28 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))
+
+
+@dataclasses.dataclass
+class D:
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+D(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    D(1, jnp.array(1.0))

test/import_hook_tester_transitive/__init__.py

@@ -17,4 +17,4 @@
 # 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.
 
-from . import another_file
+from . import another_file  # noqa: F401

test/import_hook_tester_transitive/another_file.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 equinox as eqx
 import jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+class M(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+M(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    M(1, jnp.array(1.0))

test/import_hook_tester_typeguard.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 jax.numpy as jnp
 import pytest
 
@@ -32,3 +33,31 @@ def g(x: Float32[jnp.ndarray, " b"]):
 g(jnp.array([1.0]))
 with pytest.raises(ParamError):
     g(jnp.array(1))
+
+
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/import_hook_tester_typeguard_old.py

@@ -0,0 +1,63 @@
+# 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 jax.numpy as jnp
+import pytest
+
+from jaxtyping import Float32
+
+from .helpers import ParamError
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+# Typeguard 3.0 no longer supports this.
+#
+# class M(eqx.Module):
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# M(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     M(1, jnp.array(1.0))
+#
+#
+#
+# @dataclasses.dataclass
+# class D:
+#     foo: int
+#     bar: Float32[jnp.ndarray, " a"]
+#
+#
+# D(1, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1.0, jnp.array([1.0]))
+# with pytest.raises(ParamError):
+#     D(1, jnp.array(1.0))

test/requirements.txt

@@ -2,3 +2,4 @@ equinox>=0.5.3
 pytest>=7.0.1
 beartype>=0.10.4
 typeguard>=2.13.3
+cloudpickle>=2.2.1

test/test_array.py

@@ -17,11 +17,16 @@
 # 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 sys
+from typing import get_args, get_origin, Union
+
 import jax.numpy as jnp
 import jax.random as jr
+import numpy as np
 import pytest
+import torch
 
-from jaxtyping import AbstractDtype, Array, Float, Float32, jaxtyped, Shaped
+from jaxtyping import AbstractDtype, Array, ArrayLike, Float, Float32, jaxtyped, Shaped
 
 from .helpers import ParamError, ReturnError
 
@@ -409,3 +414,101 @@ def test_incomplete_symbolic(typecheck, getkey):
     x = jr.normal(getkey(), (4,))
     with pytest.raises(NameError):
         foo(x)
+
+
+def test_arraylike(typecheck, getkey):
+    floatlike1 = Float32[ArrayLike, ""]
+    floatlike2 = Float[ArrayLike, ""]
+    floatlike3 = Float32[ArrayLike, "4"]
+
+    assert get_origin(floatlike1) is Union
+    assert get_origin(floatlike2) is Union
+    assert get_origin(floatlike3) is Union
+    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, ""]
+    shaped2 = Shaped[ArrayLike, "4"]
+    assert get_origin(shaped1) is Union
+    assert get_origin(shaped2) is Union
+    assert set(get_args(shaped1)) == {
+        Shaped[Array, ""],
+        Shaped[np.ndarray, ""],
+        Shaped[np.bool_, ""],
+        Shaped[np.number, ""],
+        bool,
+        int,
+        float,
+        complex,
+    }
+    assert set(get_args(shaped2)) == {
+        Shaped[Array, "4"],
+        Shaped[np.ndarray, "4"],
+        Shaped[np.bool_, "4"],
+        Shaped[np.number, "4"],
+    }
+
+
+def test_subclass():
+    assert issubclass(Float[Array, ""], Array)
+    assert issubclass(Float[np.ndarray, ""], np.ndarray)
+    assert issubclass(Float[torch.Tensor, ""], torch.Tensor)
+
+
+def test_ignored_names():
+    x = Float[np.ndarray, "foo=4"]
+
+    assert isinstance(np.zeros(4), x)
+    assert not isinstance(np.zeros(5), x)
+    assert not isinstance(np.zeros((4, 5)), x)
+
+    y = Float[np.ndarray, "bar qux foo=bar+qux"]
+
+    assert isinstance(np.zeros((2, 3, 5)), y)
+    assert not isinstance(np.zeros((2, 3, 6)), y)
+
+    z = Float[np.ndarray, "bar #foo=bar"]
+
+    assert isinstance(np.zeros((3, 3)), z)
+    assert isinstance(np.zeros((3, 1)), z)
+    assert not isinstance(np.zeros((3, 4)), z)
+
+    # Weird but legal
+    w = Float[np.ndarray, "bar foo=#bar"]
+
+    assert isinstance(np.zeros((3, 3)), w)
+    assert isinstance(np.zeros((3, 1)), w)
+    assert not isinstance(np.zeros((3, 4)), w)
+
+
+def test_symbolic_functions():
+    x = Float[np.ndarray, "foo bar min(foo,bar)"]
+
+    assert isinstance(np.zeros((2, 3, 2)), x)
+    assert isinstance(np.zeros((3, 2, 2)), x)
+    assert not isinstance(np.zeros((3, 2, 4)), x)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10")
+def test_py310_unions():
+    x = np.zeros(3)
+    y = Shaped[Array | np.ndarray, "_"]
+    assert isinstance(x, get_args(y))

test/test_decorator.py

@@ -0,0 +1,77 @@
+import abc
+
+from jaxtyping import jaxtyped
+
+
+class M(metaclass=abc.ABCMeta):
+    @jaxtyped
+    def f(self):
+        ...
+
+    @jaxtyped
+    @classmethod
+    def g1(cls):
+        return 3
+
+    @classmethod
+    @jaxtyped
+    def g2(cls):
+        return 4
+
+    @jaxtyped
+    @staticmethod
+    def h1():
+        return 3
+
+    @staticmethod
+    @jaxtyped
+    def h2():
+        return 4
+
+    @jaxtyped
+    @abc.abstractmethod
+    def i1(self):
+        ...
+
+    @abc.abstractmethod
+    @jaxtyped
+    def i2(self):
+        ...
+
+
+class N:
+    @jaxtyped
+    @property
+    def j1(self):
+        return 3
+
+    @property
+    @jaxtyped
+    def j2(self):
+        return 4
+
+
+def test_identity():
+    assert M.f is M.f
+
+
+def test_classmethod():
+    assert M.g1() == 3
+    assert M.g2() == 4
+
+
+def test_staticmethod():
+    assert M.h1() == 3
+    assert M.h2() == 4
+
+
+# Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
+# @abstractmethod
+def test_abstractmethod():
+    assert M.i1.__isabstractmethod__
+    assert M.i2.__isabstractmethod__
+
+
+def test_property():
+    assert N().j1 == 3
+    assert N().j2 == 4

test/test_import_hook.py

@@ -22,13 +22,33 @@ import pytest
 from jaxtyping import install_import_hook
 
 
+def test_import_hook_typeguard_old():
+    hook = install_import_hook(
+        "test.import_hook_tester_typeguard_old", ("typeguard", "typechecked")
+    )
+    with hook:
+        from . import import_hook_tester_typeguard_old  # noqa: F401
+
+
 def test_import_hook_typeguard():
     hook = install_import_hook(
-        "test.import_hook_tester_typeguard", ("typeguard", "typechecked")
+        "test.import_hook_tester_typeguard", "typeguard.typechecked"
     )
-    from . import import_hook_tester_typeguard  # noqa: F401
+    with hook:
+        from . import import_hook_tester_typeguard  # noqa: F401
 
-    hook.uninstall()
+
+def test_import_hook_beartype_old():
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        pytest.skip("Beartype not installed")
+    else:
+        hook = install_import_hook(
+            "test.import_hook_tester_beartype_old", ("beartype", "beartype")
+        )
+        with hook:
+            from . import import_hook_tester_beartype_old  # noqa: F401
 
 
 def test_import_hook_beartype():
@@ -38,26 +58,35 @@ def test_import_hook_beartype():
         pytest.skip("Beartype not installed")
     else:
         hook = install_import_hook(
-            "test.import_hook_tester_beartype", ("beartype", "beartype")
+            "test.import_hook_tester_beartype", "beartype.beartype"
         )
-        from . import import_hook_tester_beartype  # noqa: F401
+        with hook:
+            from . import import_hook_tester_beartype  # noqa: F401
 
-        hook.uninstall()
+
+def test_import_hook_beartype_full():
+    try:
+        import beartype  # noqa: F401
+    except ImportError:
+        pytest.skip("Beartype not installed")
+    else:
+        bearchecker = "beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"  # noqa: E501
+        hook = install_import_hook("test.import_hook_tester_beartype_full", bearchecker)
+        with hook:
+            from . import import_hook_tester_beartype_full  # noqa: F401
 
 
 def test_import_hook_transitive():
     hook = install_import_hook(
-        "test.import_hook_tester_transitive", ("typeguard", "typechecked")
+        "test.import_hook_tester_transitive", "beartype.beartype"
     )
-    from . import import_hook_tester_transitive  # noqa: F401
-
-    hook.uninstall()
+    with hook:
+        from . import import_hook_tester_transitive  # noqa: F401
 
 
 def test_import_hook_broken_checker():
     hook = install_import_hook(
-        "test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
+        "test.import_hook_tester_broken_checker", "jaxtyping.does_not_exist"
     )
-    with pytest.raises(AttributeError):
+    with hook, pytest.raises(AttributeError):
         from . import import_hook_tester_broken_checker  # noqa: F401
-    hook.uninstall()

test/test_pytree.py

@@ -17,7 +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.
 
-from typing import Tuple, Union
+from typing import NamedTuple, Tuple, Union
 
 import equinox as eqx
 import jax
@@ -155,3 +155,33 @@ def test_pytree_tuple(typecheck):
         g([1, 1])
     with pytest.raises(ParamError):
         g([(1, 1), "hi"])
+
+
+def test_pytree_namedtuple(typecheck):
+    class CustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    class OtherCustomNamedTuple(NamedTuple):
+        x: Float[jnp.ndarray, "a b"]
+        y: Float[jnp.ndarray, "b c"]
+
+    @typecheck
+    def g(x: PyTree[CustomNamedTuple]):
+        ...
+
+    g(
+        CustomNamedTuple(
+            x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+            y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+        )
+    )
+    with pytest.raises(ParamError):
+        g(object())
+    with pytest.raises(ParamError):
+        g(
+            OtherCustomNamedTuple(
+                x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)),
+                y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
+            )
+        )

test/test_serialisation.py

@@ -0,0 +1,16 @@
+import cloudpickle
+import numpy as np
+import torch
+
+from jaxtyping import AbstractArray, Array, Shaped
+
+
+def test_pickle():
+    x = cloudpickle.dumps(Shaped[Array, ""])
+    y = cloudpickle.dumps(AbstractArray)
+    z = cloudpickle.dumps(Shaped[np.ndarray, ""])
+    w = cloudpickle.dumps(Shaped[torch.Tensor, ""])
+    cloudpickle.loads(x)
+    cloudpickle.loads(y)
+    cloudpickle.loads(z)
+    cloudpickle.loads(w)

test/test_threading.py

@@ -0,0 +1,64 @@
+# 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
+
+import jax.numpy as jnp
+from typeguard import typechecked
+
+from jaxtyping import Array, Float, jaxtyped
+
+
+class _ErrorableThread(threading.Thread):
+    def run(self):
+        try:
+            super().run()
+        except Exception as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        super().join(timeout)
+        if hasattr(self, "exc"):
+            raise self.exc
+
+
+def test_threading_jaxtyped():
+    @jaxtyped
+    @typechecked
+    def add(x: Float[Array, "a b"], y: Float[Array, "a b"]) -> Float[Array, "a b"]:
+        return x + y
+
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        b = jnp.array([[2.0, 3.0]])
+        add(a, b)
+
+    thread = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()
+
+
+def test_threading_nojaxtyped():
+    def run():
+        a = jnp.array([[1.0, 2.0]])
+        assert isinstance(a, Float[Array, "..."])
+
+    thread = _ErrorableThread(target=run)
+    thread.start()
+    thread.join()