Bump to Py3.9

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.
 

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.
 

jaxtyping/_array_types.py

@@ -23,16 +23,7 @@ import re
 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
 
@@ -108,9 +99,9 @@ _AbstractDim = Union[Literal[_anonymous_dim], _NamedDim, _FixedDim, _SymbolicDim
 
 
 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):
@@ -214,9 +205,9 @@ class _MetaAbstractArray(type):
     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):
@@ -290,8 +281,8 @@ class AbstractArray(metaclass=_MetaAbstractArray):
     """
 
     array_type: Any
-    dtypes: List[str]
-    dims: Tuple[_AbstractDimOrVariadicDim, ...]
+    dtypes: list[str]
+    dims: tuple[_AbstractDimOrVariadicDim, ...]
     index_variadic: Optional[int]
     dim_str: str
 
@@ -518,7 +509,7 @@ class _MetaAbstractDtype(type):
             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 "
@@ -571,7 +562,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
         ```
     """
 
-    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(
@@ -582,7 +573,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
     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:

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):

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

pyproject.toml

@@ -3,7 +3,7 @@ name = "jaxtyping"
 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"},