Version bump

This required quite a lot of refactoring! JAX supports virtual subclass registration (its metaclass is ABCMeta) but NumPy does not, so we have to actually subclass `np.ndarray`.
Simple stuff like __base__ hacking fails due to deallocator conflicts.

.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

@@ -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

API.md

@@ -144,13 +144,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)
@@ -177,7 +177,7 @@ The import hook will automatically decorate all functions, and the `__init__` me
 ```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
@@ -192,7 +192,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
 ```
 

jaxtyping/__init__.py

@@ -90,7 +90,9 @@ from .import_hook import install_import_hook as install_import_hook
 
 if typing.TYPE_CHECKING:
     # Set up to deliberately confuse a static type checker.
-    PyTree = getattr(typing, "foo" + "bar")
+    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
@@ -112,4 +114,4 @@ elif has_jax:
 
 del has_jax
 
-__version__ = "0.2.13"
+__version__ = "0.2.14"

jaxtyping/array_types.py

@@ -67,72 +67,24 @@ class _NamedDim:
         self.name = name
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.name != other.name:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.name, self.broadcastable))
-
 
 class _NamedVariadicDim:
     def __init__(self, name, broadcastable):
         self.name = name
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.name != other.name:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.name, self.broadcastable))
-
 
 class _FixedDim:
     def __init__(self, size, broadcastable):
         self.size = size
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.size != other.size:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.size, self.broadcastable))
-
 
 class _SymbolicDim:
     def __init__(self, expr, broadcastable):
         self.expr = expr
         self.broadcastable = broadcastable
 
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.expr != other.expr:
-            return False
-        if self.broadcastable != other.broadcastable:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.expr, self.broadcastable))
-
 
 _AbstractDimOrVariadicDim = Union[
     Literal[_anonymous_dim],
@@ -184,22 +136,6 @@ def _check_dims(
 
 
 class _MetaAbstractArray(type):
-    def __eq__(self, other):
-        if type(self) is not type(other):
-            return False
-        if self.array_type is not other.array_type:
-            return False
-        if self.dtypes != other.dtypes:
-            return False
-        if self.dims != other.dims:
-            return False
-        if self.index_variadic != other.index_variadic:
-            return False
-        return True
-
-    def __hash__(self):
-        return hash((self.array_type, self.dtypes, self.dims, self.index_variadic))
-
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
@@ -307,9 +243,17 @@ 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 False
+        return dims == (_anonymous_variadic_dim,)
     return (_any_dtype is dtypes) or any(d.startswith(dtype) for d in dtypes)
 
 
@@ -320,6 +264,186 @@ class AbstractArray(metaclass=_MetaAbstractArray):
     index_variadic: Optional[int]
 
 
+_not_made = object()
+
+
+@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:
+            # 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)
+    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(
@@ -328,8 +452,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 "
@@ -337,193 +460,17 @@ 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)
-        dims = tuple(dims)
-
-        _not_made = object()
-
-        def _make(x):
-            # 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 x is bool:
-                if _check_scalar("bool", cls.dtypes, dims):
-                    return x
-                else:
-                    return _not_made
-            elif x is int:
-                if _check_scalar("int", cls.dtypes, dims):
-                    return x
-                else:
-                    return _not_made
-            elif x is float:
-                if _check_scalar("float", cls.dtypes, dims):
-                    return x
-                else:
-                    return _not_made
-            elif x is complex:
-                if _check_scalar("complex", cls.dtypes, dims):
-                    return x
-                else:
-                    return _not_made
-            try:
-                type_str = x.__name__
-            except AttributeError:
-                type_str = repr(x)
-            if _array_name_format == "dtype_and_shape":
-                name = f"{cls.__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"
-                )
-            out = _MetaAbstractArray(
-                name,
-                (AbstractArray,),
-                dict(
-                    array_type=x,
-                    dtypes=cls.dtypes,
-                    dims=dims,
-                    index_variadic=index_variadic,
-                ),
-            )
-            if getattr(typing, "GENERATING_DOCUMENTATION", False):
-                out.__module__ = "builtins"
-            else:
-                out.__module__ = "jaxtyping"
-            return out
-
         if typing.get_origin(array_type) is typing.Union:
-            out = [_make(x) for x in typing.get_args(array_type)]
+            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:
-            out = _make(array_type)
+            out = _make_array(array_type, dim_str, cls.dtypes, cls.__name__)
+            if out is _not_made:
+                raise ValueError("Invalid jaxtyping type annotation.")
         return out
 
 

jaxtyping/decorator.py

@@ -21,17 +21,24 @@ import dataclasses
 import functools as ft
 import inspect
 import threading
+import weakref
 
 
 storage = threading.local()
 
 
+_fns = weakref.WeakKeyDictionary()
+
+
 class _Jaxtyped:
     def __init__(self, fn):
-        self.fn = fn
+        # Stored externally so that it doesn't get blatted in the `ft.wraps` below by
+        # a function that already has a `fn` attribute.
+        _fns[self] = fn
 
     def __get__(self, instance, owner):
-        return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
+        fn = _fns[self]
+        return ft.wraps(fn)(_Jaxtyped(fn.__get__(instance, owner)))
 
     def __call__(self, *args, **kwargs):
         try:
@@ -39,8 +46,9 @@ class _Jaxtyped:
         except AttributeError:
             memo_stack = storage.memo_stack = []
         memo_stack.append(({}, {}, {}))
+        fn = _fns[self]
         try:
-            return self.fn(*args, **kwargs)
+            return fn(*args, **kwargs)
         finally:
             memo_stack.pop()
 
@@ -69,8 +77,10 @@ def _jaxtyped_typechecker(typechecker):
     def _wrapper(kls):
         assert inspect.isclass(kls)
         if dataclasses.is_dataclass(kls):
-            init = jaxtyped(typechecker(kls.__init__))
-            kls.__init__ = init
+            if type(kls.__init__) is not _Jaxtyped:
+                # Extra `if` check to work around beartype bug #211
+                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,13 +65,18 @@ def _call_with_frames_removed(f, *args, **kwargs):
     return f(*args, **kwargs)
 
 
-def _optimized_cache_from_source(path, debug_override=None):
+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?
-    return cache_from_source(path, debug_override, optimization="jaxtyping4")
+    # 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):
@@ -80,6 +86,12 @@ def _dot_lookup(*elements):
     return out
 
 
+def _str_lookup(string):
+    module = ast.parse(string)
+    (expr,) = module.body
+    return expr.value
+
+
 class _JaxtypingTransformer(ast.NodeVisitor):
     def __init__(self, *, typechecker) -> None:
         self._parents: List[ast.AST] = []
@@ -96,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)])
                     )
@@ -112,7 +124,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         if self._typechecker is None:
             args = [ast.Constant(None)]
         else:
-            args = [_dot_lookup(*self._typechecker)]
+            args = [_str_lookup(self._typechecker)]
         node.decorator_list.insert(0, ast.Call(func, args, keywords=[]))
         self._parents.append(node)
         self.generic_visit(node)
@@ -137,7 +149,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
                 # Place at the end of the decorator list, as decorators
                 # frequently remove annotations from functions and we'd like to
                 # use those annotations.
-                node.decorator_list.append(_dot_lookup(*self._typechecker))
+                node.decorator_list.append(_str_lookup(self._typechecker))
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -148,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)
@@ -171,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)
 
@@ -234,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.
@@ -246,9 +259,9 @@ 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.
+        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.
@@ -286,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)

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_typeguard_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/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

@@ -23,6 +23,7 @@ import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
 import pytest
+import torch
 
 from jaxtyping import AbstractDtype, Array, ArrayLike, Float, Float32, jaxtyped, Shaped
 
@@ -463,3 +464,9 @@ def test_arraylike(typecheck, getkey):
         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)

test/test_import_hook.py

@@ -22,14 +22,35 @@ 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"
     )
     with hook:
         from . import import_hook_tester_typeguard  # noqa: F401
 
 
+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():
     try:
         import beartype  # noqa: F401
@@ -37,15 +58,27 @@ 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"
         )
         with hook:
             from . import import_hook_tester_beartype  # noqa: F401
 
 
+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", "typeguard.typechecked"
     )
     with hook:
         from . import import_hook_tester_transitive  # noqa: F401
@@ -53,7 +86,7 @@ def test_import_hook_transitive():
 
 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 hook, pytest.raises(AttributeError):
         from . import import_hook_tester_broken_checker  # noqa: F401

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

@@ -31,8 +31,6 @@ class _ErrorableThread(threading.Thread):
             super().run()
         except Exception as e:
             self.exc = e
-        finally:
-            del self._target, self._args, self._kwargs
 
     def join(self, timeout=None):
         super().join(timeout)