Bump version

Previously, using the import hook with dataclasses resulted in the `__init__` method of the dataclass being checked.
This was undesirable when using `eqx.field(converter=...)`, as the annotation didn't necessarily reflect the argument type.
A typical example was
```python
class Foo(eqx.Module):
    x: jax.Array = eqx.field(converter=jnp.ndarray)

Foo(1)  # 1 is not an array! But this code is valid.
```

After this change, we instead monkey-patch our checks to happen at the end of the `__init__` of the dataclass -- after conversion has run.

Note that this requires https://github.com/patrick-kidger/equinox/pull/524. Otherwise, Equinox does conversion too late (in `_ModuleMeta.__call__`, after `__init__` has been run).

.pre-commit-config.yaml

@@ -19,10 +19,11 @@
 
 repos:
   - repo: https://github.com/ambv/black
-    rev: 22.3.0
+    rev: 23.9.1
     hooks:
     - id: black
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: 'v0.0.255'
+    rev: 'v0.0.291'
     hooks:
       - id: ruff
+        args: ["--fix"]

README.md

@@ -49,14 +49,22 @@ Available at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxty
 
 [Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
 
-[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.
 
-[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+[Lineax](https://github.com/google/lineax): linear solvers.
+
+[BlackJAX](https://github.com/blackjax-devs/blackjax): probabilistic+Bayesian sampling.
+
+[Orbax](https://github.com/google/orbax): checkpointing (async/multi-host/multi-device).
 
 [sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
 
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
 [Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).
 
+[PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)
+
 ### Disclaimer
 
 This is not an official Google product.

jaxtyping/__init__.py

@@ -19,6 +19,7 @@
 
 import importlib.metadata
 import typing
+import warnings
 
 # First import some things as normal
 from ._array_types import (
@@ -196,4 +197,20 @@ elif has_jax:
 del has_jax
 
 
+check_equinox_version = True  # easy-to-replace line with copybara
+if check_equinox_version:
+    try:
+        eqx_version = importlib.metadata.version("equinox")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+    else:
+        major, minor, patch = eqx_version.split(".")
+        equinox_version = (int(major), int(minor), int(patch))
+        if equinox_version < (0, 11, 0):
+            warnings.warn(
+                "jaxtyping version >=0.2.23 should be used with Equinox version "
+                ">=0.11.1"
+            )
+
+
 __version__ = importlib.metadata.version("jaxtyping")

jaxtyping/_array_types.py

@@ -144,10 +144,19 @@ def _check_dims(
 def _is_jax_extended_dtype(dtype: Any) -> bool:
     if not has_jax:
         return False
-    if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
-        return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
-    else:  # jax<=0.4.13
-        return jax.core.is_opaque_dtype(dtype)
+    try:
+        is_dtype = issubclass(dtype, jax.numpy.generic)
+    except TypeError:
+        # `dtype` not a class
+        return False
+    else:
+        if is_dtype:
+            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
+                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
+            else:  # jax<=0.4.13
+                return jax.core.is_opaque_dtype(dtype)
+        else:
+            return False
 
 
 class _MetaAbstractArray(type):

jaxtyping/_decorator.py

@@ -23,6 +23,7 @@ import inspect
 import threading
 import types
 import weakref
+from typing import get_args, get_origin
 
 
 try:
@@ -72,7 +73,7 @@ def jaxtyped(fn):
     then the old one is returned to.
 
     For example, this means you could leave off the `@jaxtyped` decorator to enforce
-    that this function use the same axes sizes as the function it was called from.
+    that this function use the same axis sizes as the function it was called from.
 
     Likewise, this means you can use `isinstance` checks inside a function body
     and have them contribute to the same collection of consistency checks performed
@@ -134,7 +135,59 @@ def jaxtyped(fn):
         return wrapped_fn
 
 
+@jaxtyped
+def _check_dataclass_annotations(self, typechecker):
+    for field in dataclasses.fields(self):
+        for kls in self.__class__.__mro__:
+            try:
+                annotation = kls.__annotations__[field.name]
+            except KeyError:
+                pass
+            else:
+                break
+        else:
+            raise TypeError
+        if isinstance(annotation, str):
+            # Don't support stringified annotations. These are basically impossible to
+            # resolve correctly, so just skip them.
+            # This does mean that annotations like `type["Foo"]` will just fail. There
+            # doesn't seem to be any way to even detect a partially-stringified
+            # annotation.
+            continue
+        if get_origin(annotation) is type:
+            args = get_args(annotation)
+            if len(args) == 1 and isinstance(args[0], str):
+                # We also special-case this one kind of partially-stringified type
+                # annotation, so as to support Equinox <v0.11.1.
+                # This was fixed in Equinox in
+                # https://github.com/patrick-kidger/equinox/pull/543
+                continue
+        try:
+            value = getattr(self, field.name)
+        except AttributeError:
+            continue  # allow uninitialised fields, which are allowed on dataclasses
+
+        @typechecker
+        def typecheck(x: annotation):
+            pass
+
+        typecheck(value)
+
+
 def _jaxtyped_typechecker(typechecker):
+    """A decorator added by the import hook to all classes. Only affects dataclasses.
+
+    Will be called as
+    ```
+    @_jaxtyped_typechecker(beartype.beartype)
+    @dataclasses.dataclass
+    class SomeDataclass:
+        ...
+    ```
+
+    After initialisation, this will check that all fields of the dataclass match their
+    specified type annotation.
+    """
     # typechecker is expected to probably be either `typeguard.typechecked`, or
     # `beartype.beartype`, or `None`.
 
@@ -144,8 +197,32 @@ def _jaxtyped_typechecker(typechecker):
     def _wrapper(kls):
         assert inspect.isclass(kls)
         if dataclasses.is_dataclass(kls):
-            init = jaxtyped(typechecker(kls.__init__))
-            kls.__init__ = init
+            # This does not check that the arguments passed to `__init__` match the
+            # type annotations. There may be a custom user `__init__`, or a
+            # dataclass-generated `__init__` used alongside
+            # `equinox.field(converter=...)`
+
+            init = kls.__init__
+
+            @ft.wraps(init)
+            def __init__(self, *args, **kwargs):
+                init(self, *args, **kwargs)
+                # `kls.__init__` is late-binding to the `__init__` function that we're
+                # in now. (Or to someone else's monkey-patch.) Either way, this checks
+                # that we're in the "top-level" `__init__`, and not one that is being
+                # called via `super()`. We don't want to trigger too early, before all
+                # fields have been assigned.
+                #
+                # We're not checking `if self.__class__ is kls` because Equinox replaces
+                # the with a defrozen version of itself during `__init__`, so the check
+                # wouldn't trigger.
+                #
+                # We're not doing this check by adding it to the end of the metaclass
+                # `__call__`, because Python doesn't allow you monkey-patch metaclasses.
+                if self.__class__.__init__ is kls.__init__:
+                    _check_dataclass_annotations(self, typechecker)
+
+            kls.__init__ = __init__
         return kls
 
     return _wrapper

jaxtyping/_import_hook.py

@@ -76,8 +76,10 @@ def _optimized_cache_from_source(typechecker_hash, /, path, debug_override=None)
     # 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.
+    # Version 7: Using the same md5 hash of the `typechecker` argument
+    #    for importlib and decorator lookup.
     return cache_from_source(
-        path, debug_override, optimization=f"jaxtyping6{typechecker_hash}"
+        path, debug_override, optimization=f"jaxtyping7{typechecker_hash}"
     )
 
 
@@ -88,20 +90,63 @@ def _dot_lookup(*elements):
     return out
 
 
-def _str_lookup(string):
-    module = ast.parse(string)
-    (expr,) = module.body
-    return expr.value
+class Typechecker:
+    lookup = {}
+
+    def __init__(self, typechecker):
+        self.ast = None
+
+        if isinstance(typechecker, str):
+            # If the typechecker is a string, then we parse it
+            string_to_eval = (
+                "def f(x, *args, **kwargs):\n"
+                + f"  import {typechecker.split('.', 1)[0]}\n"
+                + f"  return {typechecker}(x, *args, **kwargs)"
+            )
+
+            # md5 hashing instead of __hash__
+            # because __hash__ is different for each Python session
+            self.hash = hashlib.md5(typechecker.encode("utf-8")).hexdigest()
+
+            vars = {}
+            exec(string_to_eval, {}, vars)
+            Typechecker.lookup[self.hash] = vars["f"]
+
+        elif typechecker is None:
+            # If it is None, ignore it silently (use dummy decorator)
+            self.hash = 0
+            Typechecker.lookup[self.hash] = lambda x, *_, **__: x
+        else:
+            # Passed typechecker is invalid
+            raise TypeError(
+                "Jaxtyping typechecker has to be either a string or a None."
+            )
+
+    def get_hash(self):
+        return self.hash
+
+    def get_ast(self):
+        # we compile AST only if we missed importlib cache
+        if self.ast is None:
+            self.ast = (
+                ast.parse(
+                    f"@jaxtyping._import_hook.Typechecker.lookup['{self.hash}']\n"
+                    "def _():\n    ..."
+                )
+                .body[0]
+                .decorator_list[0]
+            )
+
+        return self.ast
 
 
-class _JaxtypingTransformer(ast.NodeVisitor):
-    def __init__(self, *, typechecker) -> None:
+class JaxtypingTransformer(ast.NodeVisitor):
+    def __init__(self, *, typechecker: Typechecker) -> None:
         self._parents: list[ast.AST] = []
         self._typechecker = typechecker
 
     def visit_Module(self, node: ast.Module):
-        # Insert "import typeguard; import jaxtping" after any "from __future__ ..."
-        # imports
+        # Insert "import jaxtyping" after any "from __future__ ..." imports
         for i, child in enumerate(node.body):
             if isinstance(child, ast.ImportFrom) and child.module == "__future__":
                 continue
@@ -109,11 +154,6 @@ class _JaxtypingTransformer(ast.NodeVisitor):
                 continue  # module docstring
             else:
                 node.body.insert(i, ast.Import(names=[ast.alias("jaxtyping", None)]))
-                if self._typechecker is not None:
-                    typechecker_module, _ = self._typechecker.split(".", 1)
-                    node.body.insert(
-                        i, ast.Import(names=[ast.alias(typechecker_module, None)])
-                    )
                 break
 
         self._parents.append(node)
@@ -123,11 +163,9 @@ class _JaxtypingTransformer(ast.NodeVisitor):
 
     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=[]))
+        node.decorator_list.insert(
+            0, ast.Call(func, [self._typechecker.get_ast()], keywords=[])
+        )
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -151,11 +189,11 @@ class _JaxtypingTransformer(ast.NodeVisitor):
             # FWIW, typeguard also wants to be at the end of the decorator list, as it
             # works by recompiling the wrapped function.
             node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
-            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.
-                node.decorator_list.append(_str_lookup(self._typechecker))
+            # Place typechecker at the end of the decorator list, as decorators
+            # frequently remove annotations from functions and we'd like to
+            # use those annotations.
+            node.decorator_list.append(self._typechecker.get_ast())
+
         self._parents.append(node)
         self.generic_visit(node)
         self._parents.pop()
@@ -163,12 +201,9 @@ class _JaxtypingTransformer(ast.NodeVisitor):
 
 
 class _JaxtypingLoader(SourceFileLoader):
-    def __init__(self, *args, typechecker, **kwargs):
+    def __init__(self, *args, typechecker: Typechecker, **kwargs):
         super().__init__(*args, **kwargs)
         self._typechecker = typechecker
-        self._typechecker_hash = hashlib.md5(
-            self._typechecker.encode("utf-8")
-        ).hexdigest()
 
     def source_to_code(self, data, path, *, _optimize=-1):
         source = decode_source(data)
@@ -181,7 +216,7 @@ class _JaxtypingLoader(SourceFileLoader):
             dont_inherit=True,
             optimize=_optimize,
         )
-        tree = _JaxtypingTransformer(typechecker=self._typechecker).visit(tree)
+        tree = JaxtypingTransformer(typechecker=self._typechecker).visit(tree)
         ast.fix_missing_locations(tree)
         return _call_with_frames_removed(
             compile, tree, path, "exec", dont_inherit=True, optimize=_optimize
@@ -192,7 +227,7 @@ class _JaxtypingLoader(SourceFileLoader):
         # patch safe
         with patch(
             "importlib._bootstrap_external.cache_from_source",
-            ft.partial(_optimized_cache_from_source, self._typechecker_hash),
+            ft.partial(_optimized_cache_from_source, self._typechecker.get_hash()),
         ):
             return super().exec_module(module)
 
@@ -204,7 +239,7 @@ class _JaxtypingFinder(MetaPathFinder):
     Should not be used directly, but rather via `install_import_hook`.
     """
 
-    def __init__(self, modules, original_pathfinder, typechecker):
+    def __init__(self, modules, original_pathfinder, typechecker: Typechecker):
         self.modules = modules
         self._original_pathfinder = original_pathfinder
         self._typechecker = typechecker
@@ -291,8 +326,8 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
         install_import_hook(["foo", "bar.baz"], ...)
         ```
 
-    The import hook will automatically decorate all functions, and the `__init__` method
-    of dataclasses.
+    The import hook will automatically decorate all functions, and check the attributes
+    assigned to dataclasses.
 
     If the function already has any decorators on it, then both the `@jaxtyped` and the
     typechecker decorators will get added at the bottom of the decorator list, e.g.
@@ -366,6 +401,28 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
 
         (This is the author's preferred approach to performing runtime type-checking
         with jaxtyping!)
+
+    !!! warning
+
+        Stringified dataclass annotations, e.g.
+        ```python
+        @dataclass()
+        class Foo:
+            x: "int"
+        ```
+        will be silently skipped without checking them. This is because these are
+        essentially impossible to resolve at runtime. Such stringified annotations
+        typically occur either when using them for forward references, or when using
+        `from __future__ import annotations`. (You should never use the latter, it is
+        largely incompatible with runtime type checking.)
+
+        Partially stringified dataclass annotations, e.g.
+        ```python
+        @dataclass()
+        class Foo:
+            x: tuple["int"]
+        ```
+        will likely raise an error, and must not be used at all.
     """  # noqa: E501
 
     if isinstance(modules, str):
@@ -385,6 +442,7 @@ def install_import_hook(modules: Union[str, Sequence[str]], typechecker: Optiona
     else:
         raise RuntimeError("Cannot find a PathFinder in sys.meta_path")
 
-    hook = _JaxtypingFinder(modules, finder, typechecker)
+    wrapped_typechecker = Typechecker(typechecker)
+    hook = _JaxtypingFinder(modules, finder, wrapped_typechecker)
     sys.meta_path.insert(0, hook)
     return ImportHookManager(hook)

jaxtyping/_indirection.py

@@ -1,3 +1,22 @@
+# 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.
+
 # Note that `from typing import Annotated; Bool = Annotated`
 # does not work with static type checkers. `Annotated` is a typeform rather
 # than a type, meaning it cannot be assigned.

jaxtyping/_ipython_extension.py

@@ -1,4 +1,23 @@
-from ._import_hook import _JaxtypingTransformer
+# 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.
+
+from ._import_hook import JaxtypingTransformer, Typechecker
 
 
 try:
@@ -8,20 +27,22 @@ try:
     class ChooseTypecheckerMagics(Magics):
         @line_magic("jaxtyping.typechecker")
         def typechecker(self, typechecker):
-            # remove old _JaxtypingTransformer, if present
+            # remove old JaxtypingTransformer, if present
             self.shell.ast_transformers = list(
                 filter(
-                    lambda x: not isinstance(x, _JaxtypingTransformer),
+                    lambda x: not isinstance(x, JaxtypingTransformer),
                     self.shell.ast_transformers,
                 )
             )
 
             # add new one
             self.shell.ast_transformers.append(
-                _JaxtypingTransformer(typechecker=typechecker)
+                JaxtypingTransformer(typechecker=Typechecker(typechecker))
             )
 
-except ImportError:
+except Exception:
+    # Very broad exception-handling, as e.g. IPython will sometimes be
+    # present but fail to import for mysterious reasons.
     pass
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.22"
+version = "0.2.23"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"
@@ -23,7 +23,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 urls = {repository = "https://github.com/google/jaxtyping" }
-dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3", "typing_extensions>=3.7.4.1"]
+dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3,<3", "typing_extensions>=3.7.4.1"]
 entry-points = {pytest11 = {jaxtyping = "jaxtyping._pytest_plugin"}}
 
 [build-system]

test/conftest.py

@@ -48,3 +48,13 @@ def getkey():
         return jr.PRNGKey(random.randint(0, 2**31 - 1))
 
     return _getkey
+
+
+@pytest.fixture(scope="module")
+def beartype_or_skip():
+    yield pytest.importorskip("beartype")
+
+
+@pytest.fixture(scope="module")
+def typeguard_or_skip():
+    yield pytest.importorskip("typeguard")

test/import_hook_tester.py

@@ -0,0 +1,234 @@
+# 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 helpers import ParamError, ReturnError
+
+import jaxtyping
+from jaxtyping import Float32, Int
+
+
+#
+# Test that functions get checked
+#
+
+
+def g(x: Float32[jnp.ndarray, " b"]):
+    pass
+
+
+g(jnp.array([1.0]))
+with pytest.raises(ParamError):
+    g(jnp.array(1))
+
+
+#
+# Test that Equinox modules get checked
+#
+
+
+# Dataclass `__init__`, no converter
+class Mod1(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+
+Mod1(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod1(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod1(1, jnp.array(1.0))
+
+
+# Dataclass `__init__`, converter
+class Mod2(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=jnp.asarray)
+
+
+Mod2(1)  # This will fail unless we run typechecking after conversion
+
+
+class BadMod2(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=lambda x: x)
+
+
+with pytest.raises(ParamError):
+    BadMod2(1)
+with pytest.raises(ParamError):
+    BadMod2("asdf")
+
+
+# Custom `__init__`, no converter
+class Mod3(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+    def __init__(self, foo: str, bar: Float32[jnp.ndarray, " a"]):
+        self.foo = int(foo)
+        self.bar = bar
+
+
+Mod3("1", jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod3(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod3("1", jnp.array(1.0))
+
+
+# Custom `__init__`, converter
+class Mod4(eqx.Module):
+    a: Int[jnp.ndarray, ""] = eqx.field(converter=jnp.asarray)
+
+    def __init__(self, a: str):
+        self.a = int(a)
+
+
+Mod4("1")  # This will fail unless we run typechecking after conversion
+
+
+# Custom `__post_init__`, no converter
+class Mod5(eqx.Module):
+    foo: int
+    bar: Float32[jnp.ndarray, " a"]
+
+    def __post_init__(self):
+        pass
+
+
+Mod5(1, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod5(1.0, jnp.array([1.0]))
+with pytest.raises(ParamError):
+    Mod5(1, jnp.array(1.0))
+
+
+# Dataclass `__init__`, converter
+class Mod6(eqx.Module):
+    a: jnp.ndarray = eqx.field(converter=jnp.asarray)
+
+    def __post_init__(self):
+        pass
+
+
+Mod6(1)  # This will fail unless we run typechecking after conversion
+
+
+#
+# Test that dataclasses get checked
+#
+
+
+@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 that methods get checked
+#
+
+
+class N(eqx.Module):
+    a: jnp.ndarray
+
+    def __init__(self, foo: str):
+        self.a = jnp.array(1)
+
+    def foo(self, x: jnp.ndarray):
+        pass
+
+    def bar(self) -> jnp.ndarray:
+        return self.a
+
+
+n = N("hi")
+with pytest.raises(ParamError):
+    N(123)
+with pytest.raises(ParamError):
+    n.foo("not_an_array_either")
+bad_n = eqx.tree_at(lambda x: x.a, n, "not_an_array")
+with pytest.raises(ReturnError):
+    bad_n.bar()
+
+
+#
+# Test that we don't get called in `super()`.
+#
+
+
+called = False
+
+
+class Base(eqx.Module):
+    x: int
+
+    def __init__(self):
+        self.x = "not an int"
+        global called
+        assert not called
+        called = True
+
+
+class Derived(Base):
+    def __init__(self):
+        assert not called
+        super().__init__()
+        assert called
+        self.x = 2
+
+
+Derived()
+
+
+#
+# Test that stringified type annotations work
+
+
+class Foo:
+    pass
+
+
+class Bar(eqx.Module):
+    x: type[Foo]
+    y: "type[Foo]"
+    # Note that this is the *only* kind of partially-stringified type annotation that
+    # is supported. This is for compatibility with older Equinox versions.
+    z: type["Foo"]
+
+
+Bar(Foo, Foo, Foo)
+
+with pytest.raises(ParamError):
+    Bar(1, Foo, Foo)
+
+# Record that we've finished our checks successfully
+
+jaxtyping._test_import_hook_counter += 1

test/import_hook_tester_beartype.py

@@ -1,62 +0,0 @@
-# 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_full.py

@@ -1,62 +0,0 @@
-# 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

@@ -1,62 +0,0 @@
-# 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

@@ -1,62 +0,0 @@
-# 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_transitive/__init__.py

@@ -1,20 +0,0 @@
-# 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.
-
-from . import another_file  # noqa: F401

test/import_hook_tester_transitive/another_file.py

@@ -1,48 +0,0 @@
-# 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))

test/import_hook_tester_typeguard.py

@@ -1,63 +0,0 @@
-# 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/import_hook_tester_typeguard_old.py

@@ -1,63 +0,0 @@
-# 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

@@ -1,7 +1,8 @@
 beartype
 cloudpickle
 equinox
+IPython
 jaxlib
 pytest
+tensorflow
 typeguard<3
-IPython

test/test_import_hook.py

@@ -17,76 +17,105 @@
 # 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 importlib
+import importlib.metadata
+import pathlib
+import shutil
+import sys
+import tempfile
+
 import pytest
 
-from jaxtyping import install_import_hook
+import jaxtyping
 
 
-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
+_here = pathlib.Path(__file__).resolve().parent
 
 
-def test_import_hook_typeguard():
-    hook = install_import_hook(
-        "test.import_hook_tester_typeguard", "typeguard.typechecked"
-    )
-    with hook:
-        from . import import_hook_tester_typeguard  # noqa: F401
-
-
-def test_import_hook_beartype_old():
+try:
+    typeguard_version = importlib.metadata.version("typeguard")
+except Exception as e:
+    raise ImportError("Could not find typeguard version") from e
+else:
     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
-    except ImportError:
-        pytest.skip("Beartype not installed")
-    else:
-        hook = install_import_hook(
-            "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"
+        major, _, _ = typeguard_version.split(".")
+        major = int(major)
+    except Exception as e:
+        raise ImportError(
+            f"Unexpected typeguard version {typeguard_version}; not formatted as "
+            "`major.minor.patch`"
+        ) from e
+if major != 2:
+    raise ImportError(
+        "jaxtyping's tests required typeguard version 2. (Versions 3 and 4 are both "
+        "known to have bugs.)"
     )
-    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"
-    )
-    with hook, pytest.raises(AttributeError):
-        from . import import_hook_tester_broken_checker  # noqa: F401
+assert not hasattr(jaxtyping, "_test_import_hook_counter")
+jaxtyping._test_import_hook_counter = 0
+
+
+@pytest.fixture(scope="module")
+def importhook_tempdir():
+    with tempfile.TemporaryDirectory() as dir:
+        sys.path.append(dir)
+        dir = pathlib.Path(dir)
+        shutil.copyfile(_here / "helpers.py", dir / "helpers.py")
+        yield dir
+
+
+def _test_import_hook(importhook_tempdir, typechecker):
+    counter = jaxtyping._test_import_hook_counter
+    stem = f"import_hook_tester{counter}"
+    shutil.copyfile(_here / "import_hook_tester.py", importhook_tempdir / f"{stem}.py")
+
+    importlib.invalidate_caches()
+    with jaxtyping.install_import_hook(stem, typechecker):
+        importlib.import_module(stem)
+    assert counter + 1 == jaxtyping._test_import_hook_counter
+
+
+# Tests start below...
+
+
+def test_import_hook_typeguard(importhook_tempdir, typeguard_or_skip):
+    _test_import_hook(importhook_tempdir, "typeguard.typechecked")
+
+
+def test_import_hook_beartype(importhook_tempdir, beartype_or_skip):
+    _test_import_hook(importhook_tempdir, "beartype.beartype")
+
+
+def test_import_hook_beartype_full(importhook_tempdir, beartype_or_skip):
+    bearchecker = "beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"  # noqa: E501
+    _test_import_hook(importhook_tempdir, bearchecker)
+
+
+def test_import_hook_typeguard_old(importhook_tempdir, typeguard_or_skip):
+    _test_import_hook(importhook_tempdir, ("typeguard", "typechecked"))
+
+
+def test_import_hook_beartype_old(importhook_tempdir, beartype_or_skip):
+    _test_import_hook(importhook_tempdir, ("beartype", "beartype"))
+
+
+def test_import_hook_broken_checker(importhook_tempdir):
+    with pytest.raises(AttributeError):
+        _test_import_hook(importhook_tempdir, "jaxtyping.does_not_exist")
+
+
+def test_import_hook_transitive(importhook_tempdir, typeguard_or_skip):
+    counter = jaxtyping._test_import_hook_counter
+    transitive_name = "jaxtyping_transitive_test"
+    transitive_dir = importhook_tempdir / transitive_name
+    transitive_dir.mkdir()
+    shutil.copyfile(_here / "import_hook_tester.py", transitive_dir / "tester.py")
+    with open(transitive_dir / "__init__.py", "w") as f:
+        f.write("from . import tester")
+        f.flush()
+
+    importlib.invalidate_caches()
+    with jaxtyping.install_import_hook(transitive_name, "typeguard.typechecked"):
+        importlib.import_module(transitive_name)
+    assert counter + 1 == jaxtyping._test_import_hook_counter

test/test_ipython_extension.py

@@ -14,7 +14,7 @@ def ip(session_ip):
     session_ip.run_cell(raw_cell="import jaxtyping")
     session_ip.run_line_magic(magic_name="load_ext", line="jaxtyping")
     session_ip.run_line_magic(
-        magic_name="jaxtyping.typechecker", line="beartype.beartype"
+        magic_name="jaxtyping.typechecker", line="typeguard.typechecked"
     )
     yield session_ip
 

test/test_tf_dtype.py

@@ -0,0 +1,13 @@
+# Tensorflow dependency kept in a separate file, so that we can optionally exclude it
+# more easily.
+import tensorflow as tf
+
+from jaxtyping import UInt
+
+
+def test_tf_dtype():
+    x = tf.constant(1, dtype=tf.uint8)
+    y = tf.constant(1, dtype=tf.float32)
+    hint = UInt[tf.Tensor, "..."]
+    assert isinstance(x, hint)
+    assert not isinstance(y, hint)