Document IPython extension and version nump

.github/workflows/release.yml

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

.github/workflows/run_tests.yml

@@ -42,7 +42,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest wheel beartype equinox jaxlib cloudpickle
+          python -m pip install -r test/requirements.txt
           python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
 
       - name: Checks with pre-commit

README.md

@@ -29,7 +29,7 @@ def accepts_pytree_of_arrays(x: PyTree[Float[Array, "batch c1 c2"]]):
 pip install jaxtyping
 ```
 
-Requires Python 3.8+.
+Requires Python 3.9+.
 
 JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
 
@@ -41,15 +41,21 @@ 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).
+[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).
 
 ### Disclaimer
 

docs/api/advanced-features.md

@@ -1,13 +1,18 @@
 # Advanced features
 
-## Abstract base classes
+## Creating your own dtypes
 
 ::: jaxtyping.AbstractDtype
     selection:
         members:
             false
 
-::: jaxtyping.AbstractArray
-    selection:
-        members:
-            false
+## Introspection
+
+If you're writing your own type hint parser, then you may wish to detect if some Python object is a jaxtyping-provided type.
+
+You can check for dtypes by doing `issubclass(x, AbstractDtype)`. For example, `issubclass(Float32, AbstractDtype)` will pass.
+
+You can check for arrays by doing `issubclass(x, AbstractArray)`. Here, `AbstractArray` is the base class for all shape-and-dtype specified arrays, e.g. it's a base class for `Float32[Array, "foo"]`.
+
+You can check for pytrees by doing `issubclass(x, PyTree)`. For example, `issubclass(PyTree[int], PyTree)` will pass.

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/pytree.md

@@ -5,4 +5,10 @@
         members:
             false
 
-Note that `jaxtyping.PyTree` is only available if JAX has been installed.
+---
+
+:::jaxtyping.PyTreeDef
+
+---
+
+Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.

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

@@ -13,7 +13,7 @@ jaxtyping is a library providing type annotations **and runtime type-checking**
 pip install jaxtyping
 ```
 
-Requires Python 3.8+.
+Requires Python 3.9+.
 
 JAX is an optional dependency, required for a few JAX-specific types. If JAX is not installed then these will not be available, but you may still use jaxtyping to provide shape/dtype annotations for PyTorch/NumPy/TensorFlow/etc.
 
@@ -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

@@ -20,24 +20,17 @@
 import importlib.metadata
 import typing
 
-
-try:
-    import jax
-except ImportError:
-    has_jax = False
-else:
-    has_jax = True
-    del jax
-
 # First import some things as normal
-from .array_types import (
+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 ._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
@@ -69,70 +62,72 @@ elif has_jax:
 
 # Import our dtypes
 if typing.TYPE_CHECKING:
-    # Note that `from typing_extensions import Annotated; Bool = Annotated`
-    # does not work with static type checkers. `Annotated` is a typeform rather
-    # than a type, meaning it cannot be assigned.
-    from typing_extensions import (
-        Annotated as BFloat16,
-        Annotated as Bool,
-        Annotated as Complex,
-        Annotated as Complex64,
-        Annotated as Complex128,
-        Annotated as Float,
-        Annotated as Float16,
-        Annotated as Float32,
-        Annotated as Float64,
-        Annotated as Inexact,
-        Annotated as Int,
-        Annotated as Int8,
-        Annotated as Int16,
-        Annotated as Int32,
-        Annotated as Int64,
-        Annotated as Integer,
-        Annotated as Key,
-        Annotated as Num,
-        Annotated as Shaped,
-        Annotated as UInt,
-        Annotated as UInt8,
-        Annotated as UInt16,
-        Annotated as UInt32,
-        Annotated as UInt64,
+    # Introduce an indirection so that we can `import X as X` to make it clear that
+    # these are public.
+    from ._indirection import (
+        BFloat16 as BFloat16,
+        Bool as Bool,
+        Complex as Complex,
+        Complex64 as Complex64,
+        Complex128 as Complex128,
+        Float as Float,
+        Float16 as Float16,
+        Float32 as Float32,
+        Float64 as Float64,
+        Inexact as Inexact,
+        Int as Int,
+        Int8 as Int8,
+        Int16 as Int16,
+        Int32 as Int32,
+        Int64 as Int64,
+        Integer as Integer,
+        Key as Key,
+        Num as Num,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
     )
 else:
-    # noqas to work around ruff bug
-    from .array_types import (
-        BFloat16 as BFloat16,  # noqa: F401
-        Bool as Bool,  # noqa: F401
-        Complex as Complex,  # noqa: F401
-        Complex64 as Complex64,  # noqa: F401
-        Complex128 as Complex128,  # noqa: F401
-        Float as Float,  # noqa: F401
-        Float16 as Float16,  # noqa: F401
-        Float32 as Float32,  # noqa: F401
-        Float64 as Float64,  # noqa: F401
-        Inexact as Inexact,  # noqa: F401
-        Int as Int,  # noqa: F401
-        Int8 as Int8,  # noqa: F401
-        Int16 as Int16,  # noqa: F401
-        Int32 as Int32,  # noqa: F401
-        Int64 as Int64,  # noqa: F401
-        Integer as Integer,  # noqa: F401
-        Key as Key,  # noqa: F401
-        Num as Num,  # noqa: F401
-        Shaped as Shaped,  # noqa: F401
-        UInt as UInt,  # noqa: F401
-        UInt8 as UInt8,  # noqa: F401
-        UInt16 as UInt16,  # noqa: F401
-        UInt32 as UInt32,  # noqa: F401
-        UInt64 as UInt64,  # noqa: F401
+    from ._array_types import (
+        BFloat16 as BFloat16,
+        Bool as Bool,
+        Complex as Complex,
+        Complex64 as Complex64,
+        Complex128 as Complex128,
+        Float as Float,
+        Float16 as Float16,
+        Float32 as Float32,
+        Float64 as Float64,
+        Inexact as Inexact,
+        Int as Int,
+        Int8 as Int8,
+        Int16 as Int16,
+        Int32 as Int32,
+        Int64 as Int64,
+        Integer as Integer,
+        Num as Num,
+        Shaped as Shaped,
+        UInt as UInt,
+        UInt8 as UInt8,
+        UInt16 as UInt16,
+        UInt32 as UInt32,
+        UInt64 as UInt64,
     )
 
+    if has_jax:
+        from ._array_types import Key as Key
 
-# Now import PyTree
+
+# Now import PyTreeDef and PyTree
 if typing.TYPE_CHECKING:
-    # Set up to deliberately confuse a static type checker.
     import typing_extensions
 
+    from jax.tree_util import PyTreeDef as PyTreeDef
+
+    # Set up to deliberately confuse a static type checker.
     PyTree: typing_extensions.TypeAlias = getattr(typing, "foo" + "bar")
     # What's going on with this madness?
     #
@@ -151,16 +146,52 @@ if typing.TYPE_CHECKING:
     # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
     # annotation, which static type checkers aren't smart enough to resolve.
 elif has_jax:
-    from .pytree_type import PyTree as PyTree  # noqa: F401
+    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
+
+    from ._pytree_type import PyTree as PyTree  # noqa: F401
 
 
 # Conveniences
 if typing.TYPE_CHECKING:
-    from jax import Array as Scalar
     from jax.random import PRNGKeyArray as PRNGKeyArray
-    from jax.typing import ArrayLike as ScalarLike
+
+    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
 

jaxtyping/_array_types.py

@@ -23,25 +23,20 @@ import re
 import sys
 import types
 import typing
-from typing import (
-    Any,
-    Dict,
-    List,
-    Literal,
-    NoReturn,
-    Optional,
-    Tuple,
-    Union,
-)
+from typing import Any, Literal, NoReturn, Optional, Union
 
 import numpy as np
 
-from .decorator import storage
+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
@@ -108,9 +103,9 @@ _AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim
 
 
 def _check_dims(
-    cls_dims: List[_AbstractDim],
-    obj_shape: Tuple[int],
-    single_memo: Dict[str, int],
+    cls_dims: list[_AbstractDim],
+    obj_shape: tuple[int],
+    single_memo: dict[str, int],
 ) -> bool:
     assert len(cls_dims) == len(obj_shape)
     for cls_dim, obj_size in zip(cls_dims, obj_shape):
@@ -123,7 +118,8 @@ def _check_dims(
                 return False
         elif type(cls_dim) is _SymbolicDim:
             try:
-                eval_size = eval(cls_dim.expr, single_memo)
+                # Make a copy to avoid `__builtins__` getting added as a key.
+                eval_size = eval(cls_dim.expr, single_memo.copy())
             except NameError as e:
                 raise NameError(
                     f"Cannot process symbolic dimension '{cls_dim.expr}' as some "
@@ -145,12 +141,21 @@ def _check_dims(
     return True
 
 
+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)
+
+
 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
@@ -213,9 +218,9 @@ class _MetaAbstractArray(type):
     def _check_shape(
         cls,
         obj,
-        single_memo: Dict[str, int],
-        variadic_memo: Dict[str, Tuple[int, ...]],
-        variadic_broadcast_memo: Dict[str, List[Tuple[int, ...]]],
+        single_memo: dict[str, int],
+        variadic_memo: dict[str, tuple[int, ...]],
+        variadic_broadcast_memo: dict[str, list[tuple[int, ...]]],
     ):
         if cls.index_variadic is None:
             if obj.ndim != len(cls.dims):
@@ -289,8 +294,8 @@ class AbstractArray(metaclass=_MetaAbstractArray):
     """
 
     array_type: Any
-    dtypes: List[str]
-    dims: Tuple[_AbstractDimOrVariadicDim, ...]
+    dtypes: list[str]
+    dims: tuple[_AbstractDimOrVariadicDim, ...]
     index_variadic: Optional[int]
     dim_str: str
 
@@ -517,11 +522,12 @@ class _MetaAbstractDtype(type):
             f'`jaxtyping.{cls.__name__}[jnp.ndarray, "..."]`.'
         )
 
-    def __getitem__(cls, item: Tuple[Any, str]):
+    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
@@ -570,7 +576,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
         ```
     """
 
-    dtypes: Union[Literal[_any_dtype], List[Union[str, re.Pattern]]]
+    dtypes: Union[Literal[_any_dtype], list[Union[str, re.Pattern]]]
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
@@ -581,7 +587,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
 
-        dtypes: Union[Literal[_any_dtype], str, List[str]] = cls.dtypes
+        dtypes: Union[Literal[_any_dtype], str, list[str]] = cls.dtypes
         if isinstance(dtypes, (str, re.Pattern)):
             dtypes = (dtypes,)
         elif dtypes is not _any_dtype:
@@ -656,12 +662,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

@@ -25,6 +25,14 @@ import types
 import weakref
 
 
+try:
+    import jax._src.traceback_util as traceback_util
+except ImportError:
+    pass
+else:
+    traceback_util.register_exclusion(__file__)
+
+
 storage = threading.local()
 
 

jaxtyping/_import_hook.py

@@ -51,12 +51,14 @@
 
 import ast
 import functools as ft
+import hashlib
 import sys
+from collections.abc import Sequence
 from importlib.abc import MetaPathFinder
 from importlib.machinery import SourceFileLoader
 from importlib.util import cache_from_source, decode_source
 from inspect import isclass
-from typing import List, Optional, Sequence, Union
+from typing import Optional, Union
 from unittest.mock import patch
 
 
@@ -94,7 +96,7 @@ def _str_lookup(string):
 
 class _JaxtypingTransformer(ast.NodeVisitor):
     def __init__(self, *, typechecker) -> None:
-        self._parents: List[ast.AST] = []
+        self._parents: list[ast.AST] = []
         self._typechecker = typechecker
 
     def visit_Module(self, node: ast.Module):
@@ -120,7 +122,7 @@ class _JaxtypingTransformer(ast.NodeVisitor):
         return node
 
     def visit_ClassDef(self, node: ast.ClassDef):
-        func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
+        func = _dot_lookup("jaxtyping", "_decorator", "_jaxtyped_typechecker")
         if self._typechecker is None:
             args = [ast.Constant(None)]
         else:
@@ -164,7 +166,9 @@ class _JaxtypingLoader(SourceFileLoader):
     def __init__(self, *args, typechecker, **kwargs):
         super().__init__(*args, **kwargs)
         self._typechecker = typechecker
-        self._typechecker_hash = str(abs(hash(self._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)

jaxtyping/_indirection.py

@@ -0,0 +1,32 @@
+# 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.
+from typing import (
+    Annotated as BFloat16,  # noqa: F401
+    Annotated as Bool,  # noqa: F401
+    Annotated as Complex,  # noqa: F401
+    Annotated as Complex64,  # noqa: F401
+    Annotated as Complex128,  # noqa: F401
+    Annotated as Float,  # noqa: F401
+    Annotated as Float16,  # noqa: F401
+    Annotated as Float32,  # noqa: F401
+    Annotated as Float64,  # noqa: F401
+    Annotated as Inexact,  # noqa: F401
+    Annotated as Int,  # noqa: F401
+    Annotated as Int8,  # noqa: F401
+    Annotated as Int16,  # noqa: F401
+    Annotated as Int32,  # noqa: F401
+    Annotated as Int64,  # noqa: F401
+    Annotated as Integer,  # noqa: F401
+    Annotated as Key,  # noqa: F401
+    Annotated as Num,  # noqa: F401
+    Annotated as Shaped,  # noqa: F401
+    Annotated as UInt,  # noqa: F401
+    Annotated as UInt8,  # noqa: F401
+    Annotated as UInt16,  # noqa: F401
+    Annotated as UInt32,  # noqa: F401
+    Annotated as UInt64,  # noqa: F401
+)
+
+from jax import Array as Scalar  # noqa: F401
+from jax.typing import ArrayLike as ScalarLike  # noqa: F401

jaxtyping/_ipython_extension.py

@@ -0,0 +1,35 @@
+from ._import_hook import _JaxtypingTransformer
+
+
+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)
+            )
+
+except ImportError:
+    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."
+        )

jaxtyping/_pytest_plugin.py

@@ -19,7 +19,7 @@
 
 import sys
 
-from .import_hook import install_import_hook
+from ._import_hook import install_import_hook
 
 
 def pytest_addoption(parser):

jaxtyping/_pytree_type.py

@@ -35,10 +35,6 @@ class _FakePyTree(Generic[_T]):
 _FakePyTree.__name__ = "PyTree"
 _FakePyTree.__qualname__ = "PyTree"
 _FakePyTree.__module__ = "builtins"
-# Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms
-# isn't allowed.
-# Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__
-# "types", e.g. we get types.PyTree[int].
 
 
 class _MetaPyTree(type):
@@ -46,32 +42,9 @@ class _MetaPyTree(type):
         raise RuntimeError("PyTree cannot be instantiated")
 
     def __instancecheck__(cls, obj):
-        return True
+        if not hasattr(cls, "leaftype"):
+            return True  # Just `isinstance(x, PyTree)`
 
-    @ft.lru_cache(maxsize=None)
-    def __getitem__(cls, item):
-        name = str(_FakePyTree[item])
-        out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
-        if getattr(typing, "GENERATING_DOCUMENTATION", False):
-            out.__module__ = "builtins"
-        else:
-            out.__module__ = "jaxtyping"
-        return out
-
-
-try:
-    # new typeguard
-    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
-except AttributeError:
-    # old typeguard
-    _TypeCheckError = TypeError
-
-
-class _MetaSubscriptPyTree(type):
-    def __call__(self, *args, **kwargs):
-        raise RuntimeError("PyTree cannot be instantiated")
-
-    def __instancecheck__(cls, obj):
         # We could use `isinstance` here but that would fail for more complicated
         # types, e.g. PyTree[Tuple[int]]. So at least internally we make a particular
         # choice of typechecker.
@@ -93,6 +66,36 @@ class _MetaSubscriptPyTree(type):
         leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
         return all(map(is_leaftype, leaves))
 
+    # Can't return a generic (e.g. _FakePyTree[item]) because generic aliases don't do
+    # the custom __instancecheck__ that we want.
+    # We can't add that __instancecheck__  via subclassing, e.g.
+    # type("PyTree", (Generic[_T],), {}), because dynamic subclassing of typeforms
+    # isn't allowed.
+    # Likewise we can't do types.new_class("PyTree", (Generic[_T],), {}) because that
+    # has __module__ "types", e.g. we get types.PyTree[int].
+    @ft.lru_cache(maxsize=None)
+    def __getitem__(cls, item):
+        name = str(_FakePyTree[item])
+
+        class X(PyTree):
+            leaftype = item
+
+        X.__name__ = name
+        X.__qualname__ = name
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            X.__module__ = "builtins"
+        else:
+            X.__module__ = "jaxtyping"
+        return X
+
+
+try:
+    # new typeguard
+    _TypeCheckError = (TypeError, typeguard.TypeCheckError)
+except AttributeError:
+    # old typeguard
+    _TypeCheckError = TypeError
+
 
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but subclassing

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "jaxtyping"
-version = "0.2.16"
+version = "0.2.22"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
-requires-python ="~=3.8"
+requires-python ="~=3.9"
 license = {file = "LICENSE"}
 authors = [
   {name = "Patrick Kidger", email = "contact@kidger.site"},
@@ -24,7 +24,7 @@ classifiers = [
 ]
 urls = {repository = "https://github.com/google/jaxtyping" }
 dependencies = ["numpy>=1.20.0", "typeguard>=2.13.3", "typing_extensions>=3.7.4.1"]
-entry_points = {pytest11 = ["jaxtyping = jaxtyping.pytest_plugin"]}
+entry-points = {pytest11 = {jaxtyping = "jaxtyping._pytest_plugin"}}
 
 [build-system]
 requires = ["hatchling"]

test/requirements.txt

@@ -1,5 +1,7 @@
-equinox>=0.5.3
-pytest>=7.0.1
-beartype>=0.10.4
-typeguard>=2.13.3
-cloudpickle>=2.2.1
+beartype
+cloudpickle
+equinox
+jaxlib
+pytest
+typeguard<3
+IPython

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="beartype.beartype"
+    )
+    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_pytree.py

@@ -185,3 +185,11 @@ def test_pytree_namedtuple(typecheck):
                 y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)),
             )
         )
+
+
+def test_subclass_pytree():
+    x = PyTree
+    y = PyTree[int]
+    assert issubclass(x, PyTree)
+    assert issubclass(y, PyTree)
+    assert not issubclass(int, PyTree)