alias prng.threefry_2x32 in random and warn of move

jax/_src/random.py

@@ -1437,3 +1437,13 @@ def _weibull_min(key, scale, concentration, shape, dtype):
 
   # Inverse weibull CDF.
   return jnp.power(-jnp.log1p(-random_uniform), 1.0/concentration) * scale
+
+
+# TODO(frostig): remove these aliases
+
+threefry2x32_p = prng.threefry2x32_p
+
+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)

jax/random.py

@@ -105,6 +105,8 @@
   shuffle,
   split,
   t,
+  threefry_2x32,
+  threefry2x32_p,
   truncated_normal,
   uniform,
   weibull_min,