fix version issues in docs build.

typing-extensions is no longer in our requirements.txt, but is still used.
We include typing-extensions transitively through our
requirements-txt dependencies, so all GH actions are succeeding (because they
all install the test requirements). But, our docs build does not,
and is failing as a consequence.

We no longer support python 3.7, so can import from typing instead.

PiperOrigin-RevId: 549348815

haiku/_src/module.py

@@ -18,9 +18,8 @@
 import functools
 import inspect
 import re
-import sys
 from typing import (Any, Callable, ContextManager, Dict, Mapping, NamedTuple,
-                    Optional, Set, Tuple, Type, TypeVar)
+                    Optional, Protocol, Set, Tuple, Type, TypeVar)
 
 from haiku._src import base
 from haiku._src import config
@@ -29,12 +28,6 @@
 import jax
 import jax.numpy as jnp
 
-# pylint: disable=g-import-not-at-top
-if sys.version_info < (3, 8):
-  from typing_extensions import Protocol
-else:
-  from typing import Protocol
-# pylint: enable=g-import-not-at-top
 
 T = TypeVar("T")
 

haiku/_src/typing.py

@@ -16,10 +16,9 @@
 
 import abc
 import typing
-from typing import Any, Callable, Mapping, MutableMapping, Sequence
+from typing import Any, Callable, Mapping, MutableMapping, Sequence, Protocol, runtime_checkable
 
 import jax
-from typing_extensions import Protocol, runtime_checkable  # pylint: disable=multiple-statements,g-multiple-import
 
 # pytype: disable=module-attr
 try: