Repository: patrick-kidger/jaxtyping (language: python)

Issue:
jaxtyping jt.decorator.storace.memo_stack sometimes contains __builtin__ variables
Sometimes upon inspection of memo_stack i see that there is __builtins__ key containing all builtins(), i think this is not the intended the behavior. I didn't dig too deep, but this most often happens if i have arithmetic dimensions such as batch+1.

Here is a repro: 

https://colab.research.google.com/drive/17WYz1ZJuPz1NRomHYlUMDLQegBxpInnP#scrollTo=Q4qoDvog3FM0
```python 
import typeguard
import functools as ft
import jax.numpy as jnp
@typeguard.typechecked
def foo(i: jt.Int32[jt.Array, "batch"], j:jt.Int32[jt.Array, "batch+1"], ) -> jt.Int32[jt.Array, "batch"]:
  return i

@jt.jaxtyped
def bar():    
  try: 
    foo(jnp.array([1,2]), jnp.array([2, 2, 1]))
  finally:
      print(jt.decorator.storage.memo_stack)  

bar()
```
results in: 

```
[({'batch': 2, '__builtins__': {'__name__': 'builtins', '__doc__': "Built-in functions, exceptions, and other objects.\n\nNoteworthy: None is the `nil' object; Ellipsis represents `...' in slices.", '__package__': '', '__loader__': <class '_frozen_importlib.BuiltinImporter'>, '__spec__': ModuleSpec(name='builtins', loader=<class '_frozen_importlib.BuiltinImporter'>, origin='built-in'), '__build_class__': <built-in function __build_class__>, '__import__': <built-in function __import__>, 'abs': <built-in function abs>, 'all': <built-in function all>, 'any': <built-in function any>, 'ascii': <built-in function ascii>, 'bin': <built-in function bin>, 'breakpoint': <built-in function breakpoint>, 'callable': <built-in function callable>, 'chr': <built-in function chr>, 'compile': <built-in function compile>, 'delattr': <built-in function delattr>, 'dir': <built-in function dir>, 'divmod': <built-in function divmod>, 'eval': <built-in function eval>, 'exec': <built-in function exec>, 'format': <built-in function format>, 'getattr': <built-in function getattr>, 'globals': <built-in function globals>, 'hasattr': <built-in function hasattr>, 'hash': <built-in function hash>, 'hex': <built-in function hex>, 'id': <built-in function id>, 'input': <bound method Kernel.raw_input of <google.colab._kernel.Kernel object at 0x7f61db3c7cd0>>, 'isinstance': <built-in function isinstance>, 'issubclass': <built-in function issubclass>, 'iter': <built-in function iter>, 'aiter': <built-in function aiter>, 'len': <built-in function len>, 'locals': <built-in function locals>, 'max': <built-in function max>, 'min': <built-in function min>, 'next': <built-in function next>, 'anext': <built-in function anext>, 'oct': <built-in function oct>, 'ord': <built-in function ord>, 'pow': <built-in function pow>, 'print': <built-in function print>, 'repr': <built-in function repr>, 'round': <built-in function round>, 'setattr': <built-in function setattr>, 'sorted': <built-in function sorted>, 'sum': <built-in function sum>, 'vars': <built-in function vars>, 'None': None, 'Ellipsis': Ellipsis, 'NotImplemented': NotImplemented, 'False': False, 'True': True, 'bool': <class 'bool'>, 'memoryview': <class 'memoryview'>, 'bytearray': <class 'bytearray'>, 'bytes': <class 'bytes'>, 'classmethod': <class 'classmethod'>, 'complex': <class 'complex'>, 'dict': <class 'dict'>, 'enumerate': <class 'enumerate'>, 'filter': <class 'filter'>, 'float': <class 'float'>, 'frozenset': <class 'frozenset'>, 'property': <class 'property'>, 'int': <class 'int'>, 'list': <class 'list'>, 'map': <class 'map'>, 'object': <class 'object'>, 'range': <class 'range'>, 'reversed': <class 'reversed'>, 'set': <class 'set'>, 'slice': <class 'slice'>, 'staticmethod': <class 'staticmethod'>, 'str': <class 'str'>, 'super': <class 'super'>, 'tuple': <class 'tuple'>, 'type': <class 'type'>, 'zip': <class 'zip'>, '__debug__': True, 'BaseException': <class 'BaseException'>, 'Exception': <class 'Exception'>, 'TypeError': <class 'TypeError'>, 'StopAsyncIteration': <class 'StopAsyncIteration'>, 'StopIteration': <class 'StopIteration'>, 'GeneratorExit': <class 'GeneratorExit'>, 'SystemExit': <class 'SystemExit'>, 'KeyboardInterrupt': <class 'KeyboardInterrupt'>, 'ImportError': <class 'ImportError'>, 'ModuleNotFoundError': <class 'ModuleNotFoundError'>, 'OSError': <class 'OSError'>, 'EnvironmentError': <class 'OSError'>, 'IOError': <class 'OSError'>, 'EOFError': <class 'EOFError'>, 'RuntimeError': <class 'RuntimeError'>, 'RecursionError': <class 'RecursionError'>, 'NotImplementedError': <class 'NotImplementedError'>, 'NameError': <class 'NameError'>, 'UnboundLocalError': <class 'UnboundLocalError'>, 'AttributeError': <class 'AttributeError'>, 'SyntaxError': <class 'SyntaxError'>, 'IndentationError': <class 'IndentationError'>, 'TabError': <class 'TabError'>, 'LookupError': <class 'LookupError'>, 'IndexError': <class 'IndexError'>, 'KeyError': <class 'KeyError'>, 'ValueError': <class 'ValueError'>, 'UnicodeError': <class 'UnicodeError'>, 'UnicodeEncodeError': <class 'UnicodeEncodeError'>, 'UnicodeDecodeError': <class 'UnicodeDecodeError'>, 'UnicodeTranslateError': <class 'UnicodeTranslateError'>, 'AssertionError': <class 'AssertionError'>, 'ArithmeticError': <class 'ArithmeticError'>, 'FloatingPointError': <class 'FloatingPointError'>, 'OverflowError': <class 'OverflowError'>, 'ZeroDivisionError': <class 'ZeroDivisionError'>, 'SystemError': <class 'SystemError'>, 'ReferenceError': <class 'ReferenceError'>, 'MemoryError': <class 'MemoryError'>, 'BufferError': <class 'BufferError'>, 'Warning': <class 'Warning'>, 'UserWarning': <class 'UserWarning'>, 'EncodingWarning': <class 'EncodingWarning'>, 'DeprecationWarning': <class 'DeprecationWarning'>, 'PendingDeprecationWarning': <class 'PendingDeprecationWarning'>, 'SyntaxWarning': <class 'SyntaxWarning'>, 'RuntimeWarning': <class 'RuntimeWarning'>, 'FutureWarning': <class 'FutureWarning'>, 'ImportWarning': <class 'ImportWarning'>, 'UnicodeWarning': <class 'UnicodeWarning'>, 'BytesWarning': <class 'BytesWarning'>, 'ResourceWarning': <class 'ResourceWarning'>, 'ConnectionError': <class 'ConnectionError'>, 'BlockingIOError': <class 'BlockingIOError'>, 'BrokenPipeError': <class 'BrokenPipeError'>, 'ChildProcessError': <class 'ChildProcessError'>, 'ConnectionAbortedError': <class 'ConnectionAbortedError'>, 'ConnectionRefusedError': <class 'ConnectionRefusedError'>, 'ConnectionResetError': <class 'ConnectionResetError'>, 'FileExistsError': <class 'FileExistsError'>, 'FileNotFoundError': <class 'FileNotFoundError'>, 'IsADirectoryError': <class 'IsADirectoryError'>, 'NotADirectoryError': <class 'NotADirectoryError'>, 'InterruptedError': <class 'InterruptedError'>, 'PermissionError': <class 'PermissionError'>, 'ProcessLookupError': <class 'ProcessLookupError'>, 'TimeoutError': <class 'TimeoutError'>, 'open': <built-in function open>, 'copyright': Copyright (c) 2001-2023 Python Software Foundation.
All Rights Reserved.

Copyright (c) 2000 BeOpen.com.
All Rights Reserved.

Copyright (c) 1995-2001 Corporation for National Research Initiatives.
All Rights Reserved.

Copyright (c) 1991-1995 Stichting Mathematisch Centrum, Amsterdam.
All Rights Reserved., 'credits':     Thanks to CWI, CNRI, BeOpen.com, Zope Corporation and a cast of thousands
    for supporting Python development.  See [www.python.org](http://www.python.org/) for more information., 'license': Type license() to see the full license text, 'help': Type help() for interactive help, or help(object) for help about object., '__IPYTHON__': True, 'display': <function display at 0x7f61e2b089d0>, 'execfile': <function execfile at 0x7f61ca4dd240>, 'runfile': <function runfile at 0x7f61ca31add0>, '__pybind11_internals_v4_gcc_libstdcpp_cxxabi1013__': <capsule object NULL at 0x7f61ca1c1560>, 'get_ipython': <bound method InteractiveShell.get_ipython of <google.colab._shell.Shell object at 0x7f61db3c7d60>>}}, {}, {})]

```


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3eb5a0e..d42554b 100644
--- a/README.md
+++ b/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.
 
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/docs/index.md b/docs/index.md
index 44b8734..905c426 100644
--- a/docs/index.md
+++ b/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.
 
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 5080be3..1dbd5e0 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -30,14 +30,14 @@
     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 @@ class ArrayLike:
 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 @@ class ArrayLike:
         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 @@ class ArrayLike:
     )
 
     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,18 @@ class ArrayLike:
     # 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
+    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 PRNGKeyArray, Scalar, ScalarLike  # noqa: F401
 
 del has_jax
 
diff --git a/jaxtyping/array_types.py b/jaxtyping/_array_types.py
similarity index 96%
rename from jaxtyping/array_types.py
rename to jaxtyping/_array_types.py
index 45cc969..8d4120d 100644
--- a/jaxtyping/array_types.py
+++ b/jaxtyping/_array_types.py
@@ -23,20 +23,11 @@
 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 @@ def __init__(self, expr, broadcastable):
 
 
 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 "
@@ -213,9 +205,9 @@ def __instancecheck__(cls, obj):
     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 +281,8 @@ class for `Float32[Array, "foo"]`.
     """
 
     array_type: Any
-    dtypes: List[str]
-    dims: Tuple[_AbstractDimOrVariadicDim, ...]
+    dtypes: list[str]
+    dims: tuple[_AbstractDimOrVariadicDim, ...]
     index_variadic: Optional[int]
     dim_str: str
 
@@ -517,7 +509,7 @@ def __instancecheck__(cls, obj: Any) -> NoReturn:
             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 +562,7 @@ class UInt8or16(AbstractDtype):
         ```
     """
 
-    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 +573,7 @@ def __init__(self, *args, **kwargs):
     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:
diff --git a/jaxtyping/decorator.py b/jaxtyping/_decorator.py
similarity index 100%
rename from jaxtyping/decorator.py
rename to jaxtyping/_decorator.py
diff --git a/jaxtyping/import_hook.py b/jaxtyping/_import_hook.py
similarity index 98%
rename from jaxtyping/import_hook.py
rename to jaxtyping/_import_hook.py
index ab5497f..21cca35 100644
--- a/jaxtyping/import_hook.py
+++ b/jaxtyping/_import_hook.py
@@ -52,11 +52,12 @@
 import ast
 import functools as ft
 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 +95,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 +121,7 @@ def visit_Module(self, node: ast.Module):
         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:
diff --git a/jaxtyping/indirection.py b/jaxtyping/_indirection.py
similarity index 91%
rename from jaxtyping/indirection.py
rename to jaxtyping/_indirection.py
index 1b00b1e..f936282 100644
--- a/jaxtyping/indirection.py
+++ b/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
diff --git a/jaxtyping/pytree_type.py b/jaxtyping/_pytree_type.py
similarity index 81%
rename from jaxtyping/pytree_type.py
rename to 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 @@ def __call__(self, *args, **kwargs):
         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 @@ def is_leaftype(x):
         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/pyproject.toml b/pyproject.toml
index dcaf0a0..168d3ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,9 +1,9 @@
 [project]
 name = "jaxtyping"
-version = "0.2.19"
+version = "0.2.20"
 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 patch:
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/jaxtyping/pytest_plugin.py b/jaxtyping/_pytest_plugin.py
similarity index 97%
rename from jaxtyping/pytest_plugin.py
rename to jaxtyping/_pytest_plugin.py
index 1c458f1..3d9dcb6 100644
--- a/jaxtyping/pytest_plugin.py
+++ b/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):
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 g(x: PyTree[CustomNamedTuple]):
                 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)