From ff213fcce889c02b7996c896bf5533c3748a549d Mon Sep 17 00:00:00 2001
From: Roy Frostig <frostig@google.com>
Date: Sun, 15 Aug 2021 09:04:25 -0700
Subject: [PATCH] alias prng.threefry_2x32 in random and warn of move

Some call this, apparently.
---
 jax/_src/random.py | 10 ++++++++++
 jax/random.py      |  2 ++
 2 files changed, 12 insertions(+)

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 5138feaf97dc..d4ee6de42a37 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1440,3 +1440,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)
diff --git a/jax/random.py b/jax/random.py
index e65010b2e640..85569a612ef1 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -105,6 +105,8 @@
   shuffle,
   split,
   t,
+  threefry_2x32,
+  threefry2x32_p,
   truncated_normal,
   uniform,
   weibull_min,