From 3df5e9ce07d88f0fb63ed847d3dbbfd2557f3847 Mon Sep 17 00:00:00 2001
From: Sikandar <ma5161310@gmail.com>
Date: Sun, 14 Dec 2025 23:47:07 +0500
Subject: [PATCH] fix: update deprecated typing aliases to collections.abc

---
 chex/_src/asserts.py              | 3 ++-
 chex/_src/asserts_chexify_test.py | 3 ++-
 chex/_src/asserts_internal.py     | 4 ++--
 chex/_src/dataclass_test.py       | 3 ++-
 chex/_src/fake.py                 | 3 ++-
 chex/_src/pytypes.py              | 3 ++-
 chex/_src/restrict_backends.py    | 3 ++-
 chex/_src/variants.py             | 3 ++-
 8 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
index 6c08c21c..05cdc9e8 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 b4226d49..2f2aaa0b 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 2f11dd3f..d66dd763 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 f981c8a2..c2e65a0e 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 179e6f7f..f3872f62 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 c6da0777..91c812e8 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 19937d74..0123ed41 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 f067ebbb..f4c30d1d 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