From 1ff19c1d077944f578d19dc33fea0184c798a64f Mon Sep 17 00:00:00 2001
From: Nathan Tran <nathantran008@gmail.com>
Date: Sun, 16 Mar 2025 15:17:53 -0700
Subject: [PATCH] fix: update abstract type hints to use collections.abc
 instead of typing (PEP 585)

---
 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/dimensions.py           | 4 ++--
 chex/_src/fake.py                 | 3 ++-
 chex/_src/pytypes.py              | 3 ++-
 chex/_src/restrict_backends.py    | 3 ++-
 chex/_src/variants.py             | 3 ++-
 chex/_src/warnings.py             | 3 ++-
 10 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py
index 7984ddac..310f99d2 100644
--- a/chex/_src/asserts.py
+++ b/chex/_src/asserts.py
@@ -19,7 +19,8 @@
 import functools
 import inspect
 import traceback
-from typing import Any, Callable, List, Optional, Sequence, Set, Union, cast
+from collections.abc import Sequence
+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..197e75f2 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 collections.abc import Sequence
+from typing import Any, Optional, Type
 
 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 8f93caab..7a22ed90 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..29105404 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 collections.abc import Mapping
+from typing import Any, Generic, TypeVar
 import unittest
 
 from absl.testing import absltest
diff --git a/chex/_src/dimensions.py b/chex/_src/dimensions.py
index 695c5beb..f4668be6 100644
--- a/chex/_src/dimensions.py
+++ b/chex/_src/dimensions.py
@@ -14,10 +14,10 @@
 # ==============================================================================
 """Utilities to hold expected dimension sizes."""
 
-from collections.abc import Sized
+from collections.abc import Collection, Sized
 import math
 import re
-from typing import Any, Collection, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 
 Shape = Tuple[Optional[int], ...]
diff --git a/chex/_src/fake.py b/chex/_src/fake.py
index 7d3d622e..4639d784 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 collections.abc import Iterable
+from typing import Any, Callable, Optional, Union
 from unittest import mock
 from absl import flags
 import jax
diff --git a/chex/_src/pytypes.py b/chex/_src/pytypes.py
index c6da0777..ef9ad046 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 collections.abc import Iterable, Mapping, Sequence
+from typing import Any, Union
 
 import jax
 import numpy as np
diff --git a/chex/_src/restrict_backends.py b/chex/_src/restrict_backends.py
index 4487a5c4..656c10d5 100644
--- a/chex/_src/restrict_backends.py
+++ b/chex/_src/restrict_backends.py
@@ -29,7 +29,8 @@
 """
 import contextlib
 import functools
-from typing import Optional, Sequence
+from collections.abc import Sequence
+from typing import Optional
 
 # pylint: disable=g-import-not-at-top
 try:
diff --git a/chex/_src/variants.py b/chex/_src/variants.py
index f067ebbb..bd1d4991 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 collections.abc import Sequence
+from typing import Any
 import unittest
 
 from absl import flags
diff --git a/chex/_src/warnings.py b/chex/_src/warnings.py
index 28704d7c..649b88ca 100644
--- a/chex/_src/warnings.py
+++ b/chex/_src/warnings.py
@@ -15,7 +15,8 @@
 """Utilities to emit warnings."""
 
 import functools
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 import warnings