Replace deprecated typing.Hashable with collections.abc.Hashable.

docs/ext/coverage_check.py

@@ -16,7 +16,8 @@
 
 import inspect
 import types
-from typing import Any, Mapping, Sequence, Tuple
+from collections.abc import Mapping
+from typing import Any, Sequence, Tuple
 
 import optax
 from sphinx import application

examples/cifar10_resnet.ipynb

@@ -44,7 +44,8 @@
       ],
       "source": [
         "import functools\n",
-        "from typing import Any, Callable, Sequence, Tuple, Optional, Dict\n",
+        "from collections.abc import Callable\n",
+        "from typing import Any, Sequence, Tuple, Optional, Dict\n",
         "\n",
         "from flax import linen as nn\n",
         "\n",

examples/meta_learning.ipynb

@@ -45,7 +45,8 @@
       },
       "outputs": [],
       "source": [
-        "from typing import Callable, Iterator, Tuple\n",
+        "from collections.abc import Callable\n",
+        "from typing import Iterator, Tuple\n",
         "import chex\n",
         "import jax\n",
         "import jax.numpy as jnp\n",

optax/_src/alias.py

@@ -15,7 +15,8 @@
 """Aliases for popular optimizers."""
 
 import functools
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional, Union
 
 import jax.numpy as jnp
 

optax/_src/alias_test.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Tests for `alias.py`."""
 
-from typing import Any, Callable, Union
+from collections.abc import Callable
+from typing import Any, Union
 
 from absl.testing import absltest
 from absl.testing import parameterized

optax/_src/base.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Base interfaces and datatypes."""
 
-from typing import Any, Callable, NamedTuple, Optional, Protocol, runtime_checkable, Sequence, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Protocol, runtime_checkable, Sequence, Union
 
 import chex
 import jax

optax/_src/factorized.py

@@ -15,7 +15,8 @@
 """Factorized optimizers."""
 
 import dataclasses
-from typing import NamedTuple, Optional, Callable
+from collections.abc import Callable
+from typing import NamedTuple, Optional
 
 import chex
 import jax

optax/_src/linear_algebra.py

@@ -15,7 +15,8 @@
 """Linear algebra utilities used in optimisation."""
 
 import functools
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import chex
 import jax

optax/_src/linesearch.py

@@ -15,7 +15,8 @@
 """Line-searches."""
 
 import functools
-from typing import Any, Callable, NamedTuple, Optional, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Union
 
 import chex
 import jax

optax/_src/linesearch_test.py

@@ -19,7 +19,8 @@
 import io
 import itertools
 import math
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 from absl.testing import absltest
 from absl.testing import parameterized

optax/_src/utils.py

@@ -17,7 +17,8 @@
 import functools
 import inspect
 import operator
-from typing import Any, Callable, Optional, Sequence, Union
+from collections.abc import Callable
+from typing import Any, Optional, Sequence, Union
 
 import chex
 from etils import epy

optax/_src/wrappers.py

@@ -15,7 +15,7 @@
 """Transformation wrappers."""
 
 import functools
-from typing import Callable
+from collections.abc import Callable
 
 import chex
 import jax.numpy as jnp

optax/contrib/_acprop.py

@@ -18,7 +18,8 @@
 Asynchronous Update for Adaptive Gradient Methods" by Zhuang et al.
 (https://arxiv.org/abs/2110.05454).
 """
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional, Union
 
 import jax
 import jax.numpy as jnp

optax/contrib/_cocob.py

@@ -18,7 +18,8 @@
 Networks without Learning Rates Through Coin Betting" by Francesco Orabona and
 Tatiana Tommasi.
 """
-from typing import Any, Callable, NamedTuple, Optional, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Union
 
 import jax
 import jax.numpy as jnp

optax/contrib/_dog.py

@@ -22,7 +22,8 @@
   Gradient Descent Method<https://arxiv.org/pdf/2305.16284>`_, 2023.
 """
 
-from typing import Any, Callable, NamedTuple, Optional, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Union
 
 import chex
 import jax

optax/contrib/_sam.py

@@ -47,7 +47,8 @@
 """
 # pytype: skip-file
 
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 import chex
 import jax
 import jax.numpy as jnp

optax/losses/_ranking.py

@@ -47,7 +47,8 @@
 [-0.755, 0.09, 0.665]
 """
 
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 import chex
 import jax

optax/monte_carlo/control_variates.py

@@ -53,7 +53,8 @@
 
   For examples, see `control_delta_method` and `moving_avg_baseline`.
 """
-from typing import Any, Callable, Sequence
+from collections.abc import Callable
+from typing import Any, Sequence
 
 import chex
 import jax

optax/monte_carlo/stochastic_gradient_estimators.py

@@ -29,7 +29,8 @@
 """
 
 import math
-from typing import Any, Callable, Sequence
+from collections.abc import Callable
+from typing import Any, Sequence
 
 import chex
 import jax

optax/schedules/_inject.py

@@ -16,7 +16,8 @@
 
 import functools
 import inspect
-from typing import Callable, Iterable, NamedTuple, Optional, Union
+from collections.abc import Callable
+from typing import Iterable, NamedTuple, Optional, Union
 import warnings
 
 import chex

optax/transforms/_accumulation.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Gradient transformations for accumulating gradients across updates."""
 
-from typing import Any, Callable, NamedTuple, Optional, Protocol, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Protocol, Union
 
 import chex
 import jax

optax/transforms/_adding.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Additive components in gradient transformations."""
 
-from typing import Any, Callable, NamedTuple, Optional, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Optional, Union
 
 import chex
 import jax

optax/transforms/_combining.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Flexibly compose gradient transformations."""
 
-from typing import Callable, NamedTuple, Union, Mapping, Hashable
+from collections.abc import Callable, Hashable, Mapping
+from typing import NamedTuple, Union
 
 import jax
 

optax/transforms/_masking.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Wrappers that mask out part of the parameters when applying a transform."""
 
-from typing import Any, Callable, NamedTuple, Union
+from collections.abc import Callable
+from typing import Any, NamedTuple, Union
 
 import jax
 

optax/tree_utils/_random.py

@@ -14,7 +14,8 @@
 # ==============================================================================
 """Utilities to generate random pytrees."""
 
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 import chex
 import jax

optax/tree_utils/_random_test.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Tests for optax.tree_utils._random."""
 
-from typing import Callable
+from collections.abc import Callable
 
 from absl.testing import absltest
 from absl.testing import parameterized

optax/tree_utils/_state_utils.py

@@ -17,7 +17,8 @@
 import dataclasses
 import functools
 import typing
-from typing import Any, Callable, Optional, Protocol, Tuple, Union, cast
+from collections.abc import Callable
+from typing import Any, Optional, Protocol, Tuple, Union, cast
 
 import jax
 from optax._src import base