prng_key case

Niccolo-Ajroldi · Mar 10, 2022 · 10bd7fe · 10bd7fe
1 parent 877c81d
commit 10bd7fe
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/algorithmic_efficiency/spec.py b/algorithmic_efficiency/spec.py
@@ -53,7 +53,7 @@ def __init__(self, shape_tuple):
 ParameterContainer = Dict[ParameterKey, Dict[ParameterKey, Tensor]]
 ParameterTypeTree = Dict[ParameterKey, Dict[ParameterKey, ParameterType]]
 
-RandomState = Any  # Union[jax.random.prng_key, int, bytes, ...]
+RandomState = Any  # Union[jax.random.PRNGKey, int, bytes, ...]
 
 OptimizerState = Any
 Hyperparamters = Any