fix: Change deprecated jax type jax.random.KeyArray to jax.Array (#175)

adaptive-intelligent-robotics · Feb 13, 2024 · eb2cf31 · eb2cf31
1 parent a721221
commit eb2cf31
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/qdax/types.py b/qdax/types.py
@@ -45,5 +45,5 @@ def __init__(self) -> None:
 Mask: TypeAlias = jnp.ndarray
 
 # Others
-RNGKey: TypeAlias = jax.random.KeyArray
+RNGKey: TypeAlias = jax.Array
 Metrics: TypeAlias = Dict[str, jnp.ndarray]