Drop dependence on typing_extensions since typing Protocols are avail…

…able in python >= 3.8 PiperOrigin-RevId: 568377349
google · Sep 26, 2023 · cca7385 · cca7385
1 parent d839deb
commit cca7385
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 5 deletions.
diff --git a/neural_tangents/_src/predict.py b/neural_tangents/_src/predict.py
@@ -30,7 +30,7 @@
 
 import collections
 from functools import lru_cache
-from typing import Callable, Generator, Iterable, NamedTuple, Optional, Any, Union
+from typing import Callable, Generator, Iterable, NamedTuple, Optional, Any, Union, Protocol
 
 import jax
 from jax import grad
@@ -40,7 +40,6 @@
 from jax.tree_util import tree_all, tree_map
 import numpy as np
 import scipy as sp
-from typing_extensions import Protocol
 from .utils import dataclasses, utils
 from .utils.typing import Axes, Get, KernelFn
 

diff --git a/neural_tangents/_src/utils/typing.py b/neural_tangents/_src/utils/typing.py
@@ -14,8 +14,7 @@
 
 """Common Type Definitions."""
 
-from typing import Any, Generator, Optional, Sequence, TYPE_CHECKING, TypeVar, Union
-from typing_extensions import Protocol
+from typing import Any, Generator, Optional, Sequence, TYPE_CHECKING, TypeVar, Union, Protocol
 
 from jax import random
 import jax.numpy as jnp

diff --git a/setup.py b/setup.py
@@ -28,7 +28,6 @@
 INSTALL_REQUIRES = [
     'jax>=0.4.14',
     'frozendict>=2.3.8',
-    'typing_extensions>=4.7.1',
     'tf2jax>=0.3.5',
 ]