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

@@ -41,15 +41,29 @@ Available at [https://docs.kidger.site/jaxtyping](https://docs.kidger.site/jaxty
 
 ## Finally
 
-### See also: other tools in the JAX ecosystem
+### See also: other libraries in the JAX ecosystem
 
-Neural networks: [Equinox](https://github.com/patrick-kidger/equinox).
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
 
-Numerical differential equation solvers: [Diffrax](https://github.com/patrick-kidger/diffrax).
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
 
-Computer vision models: [Eqxvision](https://github.com/paganpasta/eqxvision).
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
 
-SymPy<->JAX conversion; train symbolic expressions via gradient descent: [sympy2jax](https://github.com/google/sympy2jax).
+[Optimistix](https://github.com/patrick-kidger/optimistix): root finding, minimisation, fixed points, and least squares.
+
+[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
 

docs/api/array.md

@@ -95,16 +95,15 @@ BatchImage = Float[Array, "batch channels height width"]
 
 Note that `jaxtyping.{Array, ArrayLike}` are only available if JAX has been installed.
 
-## Scalars, PRNGKeys
+## Scalars, PRNG keys
 
 For convenience, jaxtyping also includes `jaxtyping.Scalar`, `jaxtyping.ScalarLike`, and `jaxtyping.PRNGKeyArray`, defined as:
 ```python
 Scalar = Shaped[Array, ""]
 ScalarLike = Shaped[ArrayLike, ""]
 
-# Depending on the value of `JAX_ENABLE_CUSTOM_PRNG`:
-PRNGKeyArray = Key[Array, ""]
-PRNGKeyArray = UInt32[Array, "2"]
+# Left: new-style typed keys; right: old-style keys. See JEP 9263.
+PRNGKeyArray = Union[Key[Array, ""], UInt32[Array, "2"]]
 ```
 
 Recalling that shape-and-dtype specified jaxtyping arrays can be nested, this means that e.g. you can annotate the output of `jax.random.split` with `Shaped[PRNGKeyArray, "2"]`, or e.g. an integer scalar with `Int[Scalar, ""]`.

docs/api/runtime-type-checking.md

@@ -20,3 +20,15 @@ It can be a lot of effort to add `@jaxtyped` decorators all over your codebase.
 The easier option is usually to use the import hook.
 
 ::: jaxtyping.install_import_hook
+
+---
+
+#### IPython extension
+
+If you are running in an IPython environment (for example a Jupyter or Colab notebook), then the jaxtyping hook can be automatically ran via a custom magic:
+```python
+import jaxtyping
+%load_ext jaxtyping
+%jaxtyping.typechecker beartype.beartype  # or any other runtime type checker
+```
+Place this at the start of your notebook -- everything that is directly defined in the notebook, after this magic is run, will be hook'd.

docs/index.md

@@ -40,3 +40,19 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 ## Next steps
 
 Have a read of the [Array annotations](./api/array.md) documentation on the left-hand bar!
+
+## See also: other libraries in the JAX ecosystem
+
+[Equinox](https://github.com/patrick-kidger/equinox): neural networks.
+
+[Optax](https://github.com/deepmind/optax): first-order gradient (SGD, Adam, ...) optimisers.
+
+[Diffrax](https://github.com/patrick-kidger/diffrax): numerical differential equation solvers.
+
+[Lineax](https://github.com/google/lineax): linear solvers and linear least squares.
+
+[Eqxvision](https://github.com/paganpasta/eqxvision): computer vision models.
+
+[sympy2jax](https://github.com/google/sympy2jax): SymPy<->JAX conversion; train symbolic expressions via gradient descent.
+
+[Levanter](https://github.com/stanford-crfm/levanter): scalable+reliable training of foundation models (e.g. LLMs).

jaxtyping/__init__.py

@@ -19,25 +19,19 @@
 
 import importlib.metadata
 import typing
-
-
-try:
-    import jax
-except ImportError:
-    has_jax = False
-else:
-    has_jax = True
-    del jax
+import warnings
 
 # First import some things as normal
 from ._array_types import (
     AbstractArray as AbstractArray,
     AbstractDtype as AbstractDtype,
     get_array_name_format as get_array_name_format,
+    has_jax,
     set_array_name_format as set_array_name_format,
 )
 from ._decorator import jaxtyped as jaxtyped
 from ._import_hook import install_import_hook as install_import_hook
+from ._ipython_extension import load_ipython_extension as load_ipython_extension
 
 
 # Now import Array and ArrayLike
@@ -154,12 +148,28 @@ if typing.TYPE_CHECKING:
     # annotation, which static type checkers aren't smart enough to resolve.
 elif has_jax:
     if hasattr(typing, "GENERATING_DOCUMENTATION"):
+        # Most parts of the Equinox ecosystem have
+        # `typing.GENERATING_DOCUMENTATION = True` when generating documentation, to
+        # add whatever shims are necessary to get pretty docs. E.g. to have type
+        # annotations appear as just `PyTree`, not `jaxtyping.PyTree`.
+        #
+        # As jaxtyping actually wants things to appear as e.g. `jaxtyping.PyTree`,
+        # rather than just `PyTree`, then it sets
+        # `typing.GENERATING_DOCUMENTATION = False`, to disable these shims.
+        #
+        # Here we do only a `hasattr` check, as we want to get this version of
+        # `PyTreeDef` in both the jaxtyping and the Equinox(/etc.) docs.
 
         class PyTreeDef:
             """Alias for `jax.tree_util.PyTreeDef`, which is the type of the return
             from `jax.tree_util.tree_structure(...)`.
             """
 
+        if typing.GENERATING_DOCUMENTATION:
+            # Equinox etc. docs get just `PyTreeDef`.
+            # jaxtyping docs get `jaxtyping.PyTreeDef`.
+            PyTreeDef.__module__ = "builtins"
+
     else:
         from jax.tree_util import PyTreeDef as PyTreeDef
 
@@ -172,9 +182,35 @@ if typing.TYPE_CHECKING:
 
     from ._indirection import Scalar as Scalar, ScalarLike as ScalarLike
 elif has_jax:
-    from ._array_types import PRNGKeyArray, Scalar, ScalarLike  # noqa: F401
+    from ._array_types import Scalar, ScalarLike  # noqa: F401
+
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        # That is, we're generating some downstream documentation, not the jaxtyping
+        # documentation itself.
+        class PRNGKeyArray:
+            pass
+
+        PRNGKeyArray.__module__ = "builtins"
+    else:
+        from ._array_types import PRNGKeyArray
 
 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

@@ -32,7 +32,11 @@ from ._decorator import storage
 
 try:
     import jax
-except ImportError:
+except (ImportError, RuntimeError, AttributeError):
+    # We catch `RuntimeError` as JAX will throw this if it's present, but unable to run
+    # on the current machine. This fails with this error.
+    # We catch `AttributeError` as the above then leaves the module in a partially
+    # initialised state, which causes subsequent imports to fail with this error.
     has_jax = False
 else:
     has_jax = True
@@ -137,12 +141,30 @@ def _check_dims(
     return True
 
 
+def _is_jax_extended_dtype(dtype: Any) -> bool:
+    if not has_jax:
+        return False
+    try:
+        is_dtype = issubclass(dtype, jax.numpy.generic)
+    except TypeError:
+        # `dtype` not a class
+        return False
+    else:
+        if is_dtype:
+            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
+                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
+            else:  # jax<=0.4.13
+                return jax.core.is_opaque_dtype(dtype)
+        else:
+            return False
+
+
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
 
-        if has_jax and jax.core.is_opaque_dtype(obj.dtype):
+        if _is_jax_extended_dtype(obj.dtype):
             dtype = str(obj.dtype)
         elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
@@ -512,8 +534,9 @@ class _MetaAbstractDtype(type):
     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 "
-                "array type. For example `jaxtyping.Float32[jax.Array, 'foo bar']`."
+                "As of jaxtyping v0.2.0, type annotations must now include both an "
+                "array type and a shape. For example `Float[Array, 'foo bar']`.\n"
+                "Ellipsis can be used to accept any shape: `Float[Array, '...']`."
             )
         array_type, dim_str = item
         del item
@@ -648,12 +671,10 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num")
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 if has_jax:
-    if jax.config.jax_enable_custom_prng:
-        _key_regex = re.compile(r"^key<\w+>$")
-        Key = _make_dtype(_key_regex, "Key")
-        PRNGKeyArray = Key[jax.Array, ""]
-    else:
-        Key = UInt32
-        PRNGKeyArray = Key[jax.Array, "2"]
+    _key_regex = re.compile(r"^key<\w+>$")
+    Key = _make_dtype(_key_regex, "Key")
+    # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
+    # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
+    PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
     Scalar = Shaped[jax.Array, ""]
     ScalarLike = Shaped[jax.typing.ArrayLike, ""]

jaxtyping/_decorator.py

@@ -23,6 +23,15 @@ import inspect
 import threading
 import types
 import weakref
+from typing import get_args, get_origin
+
+
+try:
+    import jax._src.traceback_util as traceback_util
+except ImportError:
+    pass
+else:
+    traceback_util.register_exclusion(__file__)
 
 
 storage = threading.local()
@@ -64,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
@@ -126,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`.
 
@@ -136,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

@@ -51,6 +51,7 @@
 
 import ast
 import functools as ft
+import hashlib
 import sys
 from collections.abc import Sequence
 from importlib.abc import MetaPathFinder
@@ -75,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}"
     )
 
 
@@ -87,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
@@ -108,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)
@@ -122,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()
@@ -150,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()
@@ -162,10 +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 = str(abs(hash(self._typechecker)))
 
     def source_to_code(self, data, path, *, _optimize=-1):
         source = decode_source(data)
@@ -178,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
@@ -189,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)
 
@@ -201,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
@@ -288,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.
@@ -363,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):
@@ -382,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

@@ -0,0 +1,56 @@
+# 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:
+    from IPython.core.magic import line_magic, Magics, magics_class
+
+    @magics_class
+    class ChooseTypecheckerMagics(Magics):
+        @line_magic("jaxtyping.typechecker")
+        def typechecker(self, typechecker):
+            # remove old JaxtypingTransformer, if present
+            self.shell.ast_transformers = list(
+                filter(
+                    lambda x: not isinstance(x, JaxtypingTransformer),
+                    self.shell.ast_transformers,
+                )
+            )
+
+            # add new one
+            self.shell.ast_transformers.append(
+                JaxtypingTransformer(typechecker=Typechecker(typechecker))
+            )
+
+except Exception:
+    # Very broad exception-handling, as e.g. IPython will sometimes be
+    # present but fail to import for mysterious reasons.
+    pass
+
+
+def load_ipython_extension(ipython):
+    try:
+        ipython.register_magics(ChooseTypecheckerMagics)
+    except NameError:
+        raise NameError(
+            "ChooseTypecheckerMagics is not defined.\n\n"
+            + "You may be trying to use IPython extension without IPython installed."
+        )

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.20"
+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,6 +1,8 @@
 beartype
 cloudpickle
 equinox
+IPython
 jaxlib
 pytest
+tensorflow
 typeguard<3

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

@@ -0,0 +1,151 @@
+import pytest
+from IPython.testing.globalipapp import start_ipython
+
+from .helpers import ParamError
+
+
+@pytest.fixture(scope="session")
+def session_ip():
+    yield start_ipython()
+
+
+@pytest.fixture(scope="function")
+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="typeguard.typechecked"
+    )
+    yield session_ip
+
+
+def test_that_ipython_works(ip):
+    ip.run_cell(raw_cell="x = 1").raise_error()
+    assert ip.user_global_ns["x"] == 1
+
+
+def test_function_beartype(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(x: int):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell="f(1)").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='f("x")').raise_error()
+
+
+def test_function_none(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(a,b,c):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell='f(1,2,"k")').raise_error()
+
+
+def test_function_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_function_jaxtyped_and_jitted(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    @jax.jit
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_class_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A(eqx.Module):
+        x: Float[Array, "2"]
+
+        def do_something(self, y: Int[Array, ""]):
+            return self.x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(
+            raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))"
+        ).raise_error()
+
+
+def test_class_not_dataclass(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A:
+        def __init__(self, x):
+            self.x = x
+
+        def do_something(self, y):
+            return x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+    ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))").raise_error()

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)