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

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)