fix: changing PRNGkey in random_utils to key

mlcommons · Nov 15, 2024 · b5ad298 · b5ad298
1 parent 1005776
commit b5ad298
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
@@ -75,5 +75,5 @@ def split(seed: SeedType, num: int = 2) -> SeedType:
 def PRNGKey(seed: SeedType) -> SeedType:  # pylint: disable=invalid-name
   if FLAGS.framework == 'jax':
     _check_jax_install()
-    return jax_rng.PRNGKey(seed)
+    return jax_rng.key(seed)
   return _PRNGKey(seed)