Make typing_extensions a dev-dependency

jax/_src/api.py

@@ -28,10 +28,11 @@
 import inspect
 from typing import (Any, Callable, Generator, Hashable, Iterable, List, Literal,
                     NamedTuple, Optional, Sequence, Tuple, TypeVar, Union,
-                    overload)
+                    overload, TYPE_CHECKING)
 
 import numpy as np
-from typing_extensions import ParamSpec
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
 
 import jax
 from jax._src import linear_util as lu
@@ -108,7 +109,8 @@
 T = TypeVar("T")
 U = TypeVar("U")
 V_co = TypeVar("V_co", covariant=True)
-P = ParamSpec("P")
+if TYPE_CHECKING:
+    P = ParamSpec("P")
 
 
 map, unsafe_map = safe_map, map

jax/_src/stages.py

@@ -34,8 +34,9 @@
 
 from dataclasses import dataclass
 from typing import (Any, Dict, Generic, List, NamedTuple, Optional, Protocol,
-                    Sequence, Tuple, TypeVar)
-from typing_extensions import ParamSpec
+                    Sequence, Tuple, TypeVar, TYPE_CHECKING)
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
 
 import jax
 from jax import tree_util
@@ -620,7 +621,10 @@ def compiler_ir(self, dialect: Optional[str] = None) -> Optional[Any]:
 
 
 V_co = TypeVar("V_co", covariant=True)
-P = ParamSpec("P")
+if TYPE_CHECKING:
+    P = ParamSpec("P")
+else:
+    P = TypeVar("P")
 
 
 class Wrapped(Protocol, Generic[P, V_co]):

setup.py

@@ -67,9 +67,9 @@ def generate_proto(source):
         'numpy>=1.20',
         'opt_einsum',
         'scipy>=1.5',
-        'typing_extensions>=4.5.0',
     ],
     extras_require={
+        'dev': ['typing_extensions>=4.5.0'],
         # Minimum jaxlib version; used in testing.
         'minimum-jaxlib': [f'jaxlib=={_minimum_jaxlib_version}'],