diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index d454405..05d8662 100644 --- a/.github/workflows/run_tests.yml +++ b/.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 diff --git a/docs/api/advanced-features.md b/docs/api/advanced-features.md index f0b09c3..edb07bf 100644 --- a/docs/api/advanced-features.md +++ b/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. diff --git a/docs/api/pytree.md b/docs/api/pytree.md index 6cf7414..ff8bdf7 100644 --- a/docs/api/pytree.md +++ b/docs/api/pytree.md @@ -5,4 +5,8 @@ members: false -Note that `jaxtyping.PyTree` is only available if JAX has been installed. \ No newline at end of file +`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. diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 5080be3..22ad5a8 100644 --- a/jaxtyping/__init__.py +++ b/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 diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py index fe04bde..511c61f 100644 --- a/jaxtyping/pytree_type.py +++ b/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 diff --git a/test/requirements.txt b/test/requirements.txt index 2492cb3..99e4780 100644 --- a/test/requirements.txt +++ b/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 diff --git a/test/test_pytree.py b/test/test_pytree.py index 8144858..f4e93ab 100644 --- a/test/test_pytree.py +++ b/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)