Fixes

.isort.cfg

@@ -3,3 +3,4 @@ force_alphabetical_sort_within_sections=true
 lines_after_imports=2
 profile=black
 treat_comments_as_code=true
+extra_standard_library=typing_extensions

jaxtyping/__init__.py

@@ -51,4 +51,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.0.1"
+__version__ = "0.0.2"

jaxtyping/array_types.py

@@ -18,7 +18,8 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
-from typing import Any, Dict, List, Literal, NoReturn, Optional, Tuple, Union
+from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
+from typing_extensions import Literal
 
 import jax.numpy as jnp
 

jaxtyping/pytree_type.py

@@ -18,11 +18,27 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+from typing import Generic, TypeVar
 
 import jax
 import typeguard
 
 
+_T = TypeVar("_T")
+
+
+class _FakePyTree(Generic[_T]):
+    pass
+
+
+_FakePyTree.__name__ = "PyTree"
+_FakePyTree.__qualname__ = "PyTree"
+# 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):
     def __call__(self, *args, **kwargs):
         raise RuntimeError("PyTree cannot be instantiated")
@@ -32,7 +48,8 @@ class _MetaPyTree(type):
 
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
-        return _MetaSubscriptPyTree(f"PyTree[{item.__name__}]", (), {"leaftype": item})
+        name = str(_FakePyTree[item])
+        return _MetaSubscriptPyTree(name, (), {"leaftype": item})
 
 
 class _MetaSubscriptPyTree(type):
@@ -63,3 +80,6 @@ class _MetaSubscriptPyTree(type):
 
 
 PyTree = _MetaPyTree("PyTree", (), {})
+# Can't do `class PyTree(Generic[_T]): ...` because we need to override the
+# instancecheck for PyTree[foo], but we subclassing
+# `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.

setup.py

@@ -63,7 +63,7 @@ python_requires = "~=3.7"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.
-install_requires = ["jax>=0.3.4", "typeguard>=2.13.3"]
+install_requires = ["jax>=0.3.4", "typeguard>=2.13.3", "typing_extensions>=4.2.0"]
 
 entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])