bringing back PRNGKey instead of key, till the python311 branch is me…

…rged
mlcommons · Nov 16, 2024 · 4b2e64e · 4b2e64e
1 parent 9819868
commit 4b2e64e
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.key(seed)
+    return jax_rng.PRNGKey(seed)
   return _PRNGKey(seed)