Repository: google/jax

Issue:
vmap over PRNG keys



Discussion:

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -29,41 +29,25 @@
 from jax.lib import xla_bridge
 from jax import core
 
-class PRNGKey(object):
-  """A pseudo-random number generator (PRNG) key for use with lax.random."""
-  __slots__ = ["keypair"]
-
-  def __init__(self, seed):
-    """Create a new PRNG key.
-
-    Args:
-      seed: a scalar integer value used to initialize the PRNG key.
-
-    Returns:
-      A new PRNGKey object.
-    """
-    convert = lambda key: lax.convert_element_type(key, onp.uint32)
-    if onp.shape(seed):
-      raise TypeError("PRNGKey seed must be a scalar.")
-    if isinstance(seed, (int, onp.ndarray)):
-      # Special handling of raw integer values, which may have be 64bit even
-      # when jax_enable_x64=False and we don't want to drop the top 32 bits
-      k1 = convert(onp.bitwise_and(onp.right_shift(seed, 32), 0xFFFFFFFF))
-    else:
-      k1 = convert(lax.shift_right_logical(seed, 32))
-    k2 = convert(lax.bitwise_and(seed, 0xFFFFFFFF))
-    self.keypair = core.pack((k1, k2))
-
-  @classmethod
-  def from_keypair(cls, keypair):
-    """Internal method to create a PRNGKey instance from a raw key pair."""
-    new = cls.__new__(cls)
-    new.keypair = core.pack(keypair)
-    return new
-
-
-tree_util.register_pytree_node(PRNGKey, lambda k: (k.keypair, None),
-                               lambda _, xs: PRNGKey.from_keypair(xs))
+
+def PRNGKey(seed):
+  if onp.shape(seed):
+    raise TypeError("PRNGKey seed must be a scalar.")
+  convert = lambda k: lax.reshape(lax.convert_element_type(k, onp.uint32), [1])
+  if isinstance(seed, (int, onp.ndarray)):
+    # Special handling of raw integer values, which may have be 64bit even
+    # when jax_enable_x64=False and we don't want to drop the top 32 bits
+    k1 = convert(onp.bitwise_and(onp.right_shift(seed, 32), 0xFFFFFFFF))
+  else:
+    k1 = convert(lax.shift_right_logical(seed, 32))
+  k2 = convert(lax.bitwise_and(seed, 0xFFFFFFFF))
+  return lax.concatenate([k1, k2], 0)
+
+def is_prng_key(key):
+  try:
+    return key.shape == (2,) and key.dtype == onp.uint32
+  except AttributeError:
+    return False
 
 
 ### utilities
@@ -169,14 +153,15 @@ def split(key, num=2):
   Returns:
     A tuple of length `num` of new PRNGKey instances.
   """
-  counts = lax.tie_in(key.keypair, lax.iota(onp.uint32, num * 2))
-  bits = lax.reshape(threefry_2x32(key.keypair, counts), (num, 2))
-  keypairs = (lax.index_in_dim(bits, i, keepdims=False) for i in range(num))
-  return tuple(PRNGKey.from_keypair((kp[0], kp[1])) for kp in keypairs)
+  counts = lax.tie_in(key, lax.iota(onp.uint32, num * 2))
+  bits = lax.reshape(threefry_2x32(key, counts), (num, 2))
+  return tuple(bits)
 
 
 def _random_bits(key, bit_width, shape):
   """Sample uniform random bits of given width and shape using PRNG key."""
+  if not is_prng_key(key):
+    raise TypeError("_random_bits got invalid prng key.")
   if bit_width not in (32, 64):
     raise TypeError("requires 32- or 64-bit field width.")
   max_count = (bit_width // 32) * onp.prod(shape)
@@ -184,8 +169,8 @@ def _random_bits(key, bit_width, shape):
     # TODO(mattjj): just split the key here
     raise TypeError("requesting more random bits than a single call provides.")
 
-  counts = lax.tie_in(key.keypair, lax.iota(onp.uint32, max_count))
-  bits = threefry_2x32(key.keypair, counts)
+  counts = lax.tie_in(key, lax.iota(onp.uint32, max_count))
+  bits = threefry_2x32(key, counts)
   if bit_width == 64:
     bits = [lax.convert_element_type(x, onp.uint64) for x in np.split(bits, 2)]
     bits = (bits[0] << onp.uint64(32)) | bits[1]


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -24,6 +24,7 @@
 from jax import test_util as jtu
 from jax.abstract_arrays import ShapedArray
 from jax import lax
+from jax import random
 from jax.api import jit, grad, jvp, vjp, trace_to_jaxpr, jacfwd, jacrev, hessian
 from jax.api import vmap
 from jax.core import unit
@@ -326,6 +327,14 @@ def testDynamicSlice(self):
     expected = x[idx]
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testRandom(self):
+    seeds = vmap(random.PRNGKey)(onp.arange(10))
+    ans = vmap(partial(random.normal, shape=(3, 2)))(seeds)
+    expected = onp.stack([random.normal(random.PRNGKey(seed), (3, 2))
+                          for seed in onp.arange(10)])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert len(onp.unique(ans)) == 10 * 3 * 2
+
 
 if __name__ == '__main__':
   absltest.main()