Fix warning of the form DeprecationWarning: jax.core.Shape is depreca…

…ted. Use Shape = Sequence[int | Any] PiperOrigin-RevId: 573839993
google-deepmind · Oct 16, 2023 · 70d1302 · 70d1302
1 parent c69eeb3
commit 70d1302
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/chex/_src/pytypes.py b/chex/_src/pytypes.py
@@ -14,7 +14,7 @@
 # ==============================================================================
 """Type definitions to use for type annotations."""
 
-from typing import Any, Iterable, Mapping, Union
+from typing import Any, Iterable, Mapping, Sequence, Union
 
 import jax
 import numpy as np
@@ -52,7 +52,7 @@
 # Other types.
 Scalar = Union[float, int]
 Numeric = Union[Array, Scalar]
-Shape = jax.core.Shape
+Shape = Sequence[Union[int, Any]]
 PRNGKey = jax.Array
 PyTreeDef = jax.tree_util.PyTreeDef
 Device = jax.Device