Fix deprecated typing aliases to use collections.abc

chex/_src/asserts.py

@@ -16,10 +16,11 @@
 
 import collections
 import collections.abc
+from collections.abc import Sequence
 import functools
 import inspect
 import traceback
-from typing import Any, Callable, List, Optional, Sequence, Set, Union, cast
+from typing import Any, Callable, List, Optional, Set, Union, cast
 import unittest
 from unittest import mock
 

chex/_src/asserts_chexify_test.py

@@ -19,7 +19,8 @@
 import sys
 import threading
 import time
-from typing import Any, Optional, Sequence, Type
+from typing import Any, Optional, Type
+from collections.abc import Sequence
 
 from absl.testing import absltest
 from absl.testing import parameterized

chex/_src/asserts_internal.py

@@ -24,12 +24,12 @@
 
 import collections
 import collections.abc
-from collections.abc import Hashable
+from collections.abc import Hashable, Sequence
 import functools
 import re
 import threading
 import traceback
-from typing import Any, Sequence, Union, Callable, List, Optional, Set, Tuple, Type
+from typing import Any, Union, Callable, List, Optional, Set, Tuple, Type
 
 from absl import logging
 from chex._src import pytypes

chex/_src/dataclass_test.py

@@ -20,7 +20,8 @@
 import dataclasses
 import pickle
 import sys
-from typing import Any, Generic, Mapping, TypeVar
+from typing import Any, Generic, TypeVar
+from collections.abc import Mapping
 import unittest
 
 from absl.testing import absltest

chex/_src/fake.py

@@ -25,7 +25,8 @@
 import inspect
 import os
 import re
-from typing import Any, Callable, Iterable, Optional, Union
+from typing import Any, Callable, Optional, Union
+from collections.abc import Iterable
 from unittest import mock
 from absl import flags
 import jax

chex/_src/pytypes.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Type definitions to use for type annotations."""
 
-from typing import Any, Iterable, Mapping, Sequence, Union
+from typing import Any, Union
+from collections.abc import Iterable, Mapping, Sequence
 
 import jax
 import numpy as np

chex/_src/restrict_backends.py

@@ -29,7 +29,8 @@
 """
 import contextlib
 import functools
-from typing import Callable, Optional, Sequence
+from typing import Callable, Optional
+from collections.abc import Sequence
 
 from jax._src import compiler
 

chex/_src/variants.py

@@ -18,7 +18,8 @@
 import functools
 import inspect
 import itertools
-from typing import Any, Sequence
+from typing import Any
+from collections.abc import Sequence
 import unittest
 
 from absl import flags