alias prng.threefry_2x32 in random and warn of move

jax-ml · Aug 15, 2021 · 36479bd · 36479bd
1 parent f32a08b
commit 36479bd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/jax/_src/random.py b/jax/_src/random.py
@@ -1437,3 +1437,9 @@ def _weibull_min(key, scale, concentration, shape, dtype):
 
   # Inverse weibull CDF.
   return jnp.power(-jnp.log1p(-random_uniform), 1.0/concentration) * scale
+
+
+def threefry_2x32(keypair, count):
+  warnings.warn('jax.random.threefry_2x32 has moved to jax.prng.threefry_2x32 '
+                'and will be removed from `random` module.', FutureWarning)
+  return prng.threefry_2x32(keypair, count)
diff --git a/jax/random.py b/jax/random.py
@@ -105,6 +105,7 @@
   shuffle,
   split,
   t,
+  threefry_2x32,
   truncated_normal,
   uniform,
   weibull_min,