Repository: google/jax

Issue:
Split keys over named batch axis
A pattern that I find myself using often in jax is defining a `loss`, and then defining an outer `average_loss` that `vmap`s the inner `loss` function and then averages, i.e.
```
def loss(key, params, x):
  return ...

def average_loss(key, params, xs):
  keys = jax.random.split(key, num=xs.shape[0])
  losses =  vmap(loss, in_axes=(0, None, 0))(keys, params, xs)
  return jnp.mean(losses)
```
I can use `fold_in` and a named axis to make this cleaner:
```
def loss(key, params, x):
  key = random.fold_in(key, lax.axis_index('batch'))
  ...
  return jnp.mean(loss, 'batch')
  
average_loss = vmap(loss, in_axes=(None, None, 0), axis_name='batch')
```
But these two approaches don't give the exact same results. It would be nice to have an 'official' way to split keys over batch axes. 




Discussion:
Even if we don't dictate an official recipe, we should see if we can make these behave equivalently.
It seems like this is because `split` folds twice as many indices as `fold_in`, but it's not explained anywhere. It might have to do with the design of threefry. Any clues @mattjj @hawkinsp?
The "2x32" in threefry-2x32 refers to processing two 32-bit inputs (call them "words" for lack of a better word). Our `split` hashes `2 * n` words in order to produce `2 * n` words as `n` new keys.

Meanwhile `fold_in` is our own addition to the RNG menu. It currently seems to directly use the threefry-2x32 function as a keyed 32-bit hash. I don't know that this is any more or less valid than the alternative of, say, `fold_in(key, n) == split(key, n + 1)[n]` (ignoring the overflow case where `n` is the u32 max), but it is simple and fast.

If we want the latter meaning instead, a question then is how to efficiently compute the n'th split alone. Here's one way:

```python
from jax import numpy as jnp
from jax._src.prng import threefry_2x32, threefry_split, threefry_seed

def ith_fry(key, i, n):
  return threefry_2x32(key, jnp.array([i, n + i], dtype=jnp.uint32))

def ith_split(key, i, n):
  j1, j2 = 2 * i, 2 * i + 1
  return jnp.array([ith_fry(key, j1 % n, n)[j1 // n],
                    ith_fry(key, j2 % n, n)[j2 // n]])

def nth_split(key, n):
  return ith_split(key, n, n + 1)
```

Quick check on these:

```python
>>> key = threefry_seed(32)

>>> split = lambda key, n: jnp.array([ith_split(key, i, n) for i in range(n)])
>>> (threefry_split(key, 10) == split(key, 10)).all()
DeviceArray(True, dtype=bool)

>>> nth_split_ref = lambda key, n: threefry_split(key, n + 1)[n]
>>> all((nth_split(key, n) == nth_split_ref(key, n)).all() for n in range(50))
True
```

We likely have more room to change `split` so that the `nth_split` logic is simpler and faster. Today's `split` feeds threefry pairs of input words whose values are spaced apart by `n`, but it might be fine to give it consecutive counts instead (see `ith_fry` above). We could also rework how we reshape the threefry output in `split` as well. Today we effectively transpose the random words on the way out. That may not be necessary, and changing it would prevent the need for divmod indexing in `ith_split`. Both modifications would remove dependencies on `n`.

Anyway, we would need to do something along these lines in order to address the initial request here from @dieterichlawson. We don't only want `fold_in(key, n)` to behave as `split(key, n + 1)[n]` there. What we want is that, for all `n` and `i in range(n)`, `fold_in(key, i)` equals `split(key, n)[i]`. Because `fold_in` is independent of `n`, we'd want an `ith_split` that is independent of it as well.

Resolving patch:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -1062,13 +1062,23 @@ def threefry_2x32(keypair, count):
 
 
 def threefry_split(key: jnp.ndarray, num: int) -> jnp.ndarray:
-  return _threefry_split(key, int(num))  # type: ignore
+  if config.jax_threefry_partitionable:
+    return _threefry_split_foldlike(key, int(num))  # type: ignore
+  else:
+    return _threefry_split_original(key, int(num))  # type: ignore
 
 @partial(jit, static_argnums=(1,), inline=True)
-def _threefry_split(key, num) -> jnp.ndarray:
+def _threefry_split_original(key, num) -> jnp.ndarray:
   counts = lax.iota(np.uint32, num * 2)
   return lax.reshape(threefry_2x32(key, counts), (num, 2))
 
+@partial(jit, static_argnums=(1,), inline=True)
+def _threefry_split_foldlike(key, num) -> jnp.ndarray:
+  k1, k2 = key
+  counts1, counts2 = iota_32x2_shape((num,))
+  bits1, bits2 = threefry2x32_p.bind(k1, k2, counts1, counts2)
+  return jnp.stack([bits1, bits2], axis=1)
+
 
 def threefry_fold_in(key: jnp.ndarray, data: jnp.ndarray) -> jnp.ndarray:
   assert not data.shape
@@ -1177,7 +1187,12 @@ def _rbg_seed(seed: jnp.ndarray) -> jnp.ndarray:
   return jnp.concatenate([halfkey, halfkey])
 
 def _rbg_split(key: jnp.ndarray, num: int) -> jnp.ndarray:
-  return vmap(_threefry_split, (0, None), 1)(key.reshape(2, 2), num).reshape(num, 4)
+  if config.jax_threefry_partitionable:
+    _threefry_split = _threefry_split_foldlike
+  else:
+    _threefry_split = _threefry_split_original
+  return vmap(
+      _threefry_split, (0, None), 1)(key.reshape(2, 2), num).reshape(num, 4)
 
 def _rbg_fold_in(key: jnp.ndarray, data: jnp.ndarray) -> jnp.ndarray:
   assert not data.shape


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -367,6 +367,24 @@ def testPRNGValues(self):
         _prng_key_as_array(random.fold_in(k, 4)),
         np.array([2285895361,  433833334], dtype='uint32'))
 
+  @skipIf(not config.jax_threefry_partitionable, 'enable after upgrade')
+  def test_threefry_split_fold_in_symmetry(self):
+    with jax.default_prng_impl('threefry2x32'):
+      key = random.PRNGKey(72)
+      folds = jnp.array([random.fold_in(key, i) for i in range(8)])
+      splits = random.split(key, 8)
+      self.assertArraysEqual(folds, splits)
+
+  @skipIf(not config.jax_threefry_partitionable, 'enable after upgrade')
+  def test_threefry_split_vmapped_fold_in_symmetry(self):
+    # See https://github.com/google/jax/issues/7708
+    with jax.default_prng_impl('threefry2x32'):
+      key = random.PRNGKey(72)
+      folds = vmap(lambda k, _: random.fold_in(k, lax.axis_index('batch')),
+                   in_axes=(None, 0), axis_name='batch')(key, jnp.ones(8))
+      splits = random.split(key, 8)
+      self.assertArraysEqual(folds, splits)
+
   @jtu.sample_product([
       {"seed": 0, "type": int, "jit": True, "key": [0, 0]},
       {"seed": 0, "type": int, "jit": False, "key": [0, 0]},