chex: alias PRNGKey to jax.Array

Going forward, `jax.Array` is the correct type annotation for both new-style and old-style PRNG keys in JAX (see jax-ml/jax#17297) Note that `jax.random.KeyArray` has been aliased to `Any` under TYPE_CHECKING, and so this change will make existing annotations more strict, not less strict. PiperOrigin-RevId: 565133147
google-deepmind · Sep 13, 2023 · 6615676 · 6615676
1 parent 533aeeb
commit 6615676
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chex/_src/pytypes.py b/chex/_src/pytypes.py
@@ -51,7 +51,7 @@
 Scalar = Union[float, int]
 Numeric = Union[Array, Scalar]
 Shape = jax.core.Shape
-PRNGKey = Union[jax.random.KeyArray, jax.Array]
+PRNGKey = jax.Array
 PyTreeDef = jax.tree_util.PyTreeDef
 Device = jax.Device
 ArrayDType = type(jnp.float32)