diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
index 6c08c21..05cdc9e 100644
--- a/chex/_src/asserts.py
+++ b/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
 
diff --git a/chex/_src/asserts_chexify_test.py b/chex/_src/asserts_chexify_test.py
index b4226d4..2f2aaa0 100644
--- a/chex/_src/asserts_chexify_test.py
+++ b/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
diff --git a/chex/_src/asserts_internal.py b/chex/_src/asserts_internal.py
index 2f11dd3..d66dd76 100644
--- a/chex/_src/asserts_internal.py
+++ b/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
diff --git a/chex/_src/dataclass_test.py b/chex/_src/dataclass_test.py
index f981c8a..c2e65a0 100644
--- a/chex/_src/dataclass_test.py
+++ b/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
diff --git a/chex/_src/fake.py b/chex/_src/fake.py
index 179e6f7..f3872f6 100644
--- a/chex/_src/fake.py
+++ b/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
diff --git a/chex/_src/pytypes.py b/chex/_src/pytypes.py
index c6da077..91c812e 100644
--- a/chex/_src/pytypes.py
+++ b/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
diff --git a/chex/_src/restrict_backends.py b/chex/_src/restrict_backends.py
index 19937d7..0123ed4 100644
--- a/chex/_src/restrict_backends.py
+++ b/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
 
diff --git a/chex/_src/variants.py b/chex/_src/variants.py
index f067ebb..f4c30d1 100644
--- a/chex/_src/variants.py
+++ b/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