Now provides PyTreeDef, and can detect PyTrees via issubclass(x, PyTree)

.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

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,8 @@
         members:
             false
 
-Note that `jaxtyping.PyTree` is only available if JAX has been installed.
+`jaxtyping.PyTreeDef` is an alias for `jax.tree_util.PyTreeDef`, which is the type of the return from `jax.tree_util.tree_structure(...)`.
+
+:::jaxtyping.PyTreeDef
+
+Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.

jaxtyping/__init__.py

@@ -128,11 +128,13 @@ else:
         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,6 +153,8 @@ 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 jax.tree_util import PyTreeDef as PyTreeDef
+
     from .pytree_type import PyTree as PyTree  # noqa: F401
 
 

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

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)