Hash is now stable across python runtimes

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

@@ -30,14 +30,14 @@ else:
     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,
     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
 
 
 # Now import Array and ArrayLike
@@ -71,7 +71,7 @@ elif has_jax:
 if typing.TYPE_CHECKING:
     # Introduce an indirection so that we can `import X as X` to make it clear that
     # these are public.
-    from .indirection import (
+    from ._indirection import (
         BFloat16 as BFloat16,
         Bool as Bool,
         Complex as Complex,
@@ -98,7 +98,7 @@ if typing.TYPE_CHECKING:
         UInt64 as UInt64,
     )
 else:
-    from .array_types import (
+    from ._array_types import (
         BFloat16 as BFloat16,
         Bool as Bool,
         Complex as Complex,
@@ -125,14 +125,16 @@ else:
     )
 
     if has_jax:
-        from .array_types import Key as Key
+        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 +153,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.random import PRNGKeyArray as PRNGKeyArray
 
-    from .indirection import Scalar as Scalar, ScalarLike 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,20 +23,11 @@ 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:
@@ -108,9 +99,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 +114,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 +137,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 +214,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 +290,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,7 +518,7 @@ 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 "
@@ -570,7 +571,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 +582,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:

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

@@ -1,7 +1,7 @@
-# Note that `from typing_extensions import Annotated; Bool = Annotated`
+# 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_extensions import (
+from typing import (
     Annotated as BFloat16,  # noqa: F401
     Annotated as Bool,  # noqa: F401
     Annotated as Complex,  # noqa: F401

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.19"
+version = "0.2.21"
 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,6 @@
-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

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)