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/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/_array_types.py b/jaxtyping/_array_types.py index 42bcbf9..8d4120d 100644 --- a/jaxtyping/_array_types.py +++ b/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: diff --git a/jaxtyping/_import_hook.py b/jaxtyping/_import_hook.py index 3c91a8f..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): diff --git a/jaxtyping/_indirection.py b/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/pyproject.toml b/pyproject.toml index 076f271..168d3ae 100644 --- a/pyproject.toml +++ b/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"},