PRNGKey split patterns

[cgarciae]

On a regular training loop I usually split keys using this pattern which seems to be the most common:

key = jax.random.PRNGKey(0)

for x, y in data:
    key, step_key = jax.random.split(key, 2)
    model = train_step(step_key, model, x, y)

However, I've recently realized you can also do something like this:

key = jax.random.PRNGKey(0)

for x, y in data:
    key = jax.random.split(key, 1)[0]
    model = train_step(key, model, x, y)

Since train_step is still getting a unique key I would say this is alright but I haven't seen this used so I was just wondering if this is an anti-pattern or its ok?

[mattjj]

Thanks for the question!

The first example is better, because in it the caller guarantees random values can't be accidentally reused.

In the latter example the key is used twice: once as an argument to train_step and again as an argument to spilt on the next loop iteration. That's dangerous in the sense that the different calls to train_step might end up using the same random values. Here's an example definition of train_step where that happens:

import jax
key = jax.random.PRNGKey(0)

def train_step(key):
  key, = jax.random.split(key, 1)
  print(jax.random.normal(key, ()))
  key, = jax.random.split(key, 1)
  print(jax.random.normal(key, ()))

for i in range(2):
  print(f'iter {i}')
  key, = jax.random.split(key, 1)
  train_step(key)

iter 0
1.2683465
0.2534419
iter 1
0.2534419
0.1563971

[cgarciae]

Thanks @mattjj for the response! Indeed that would be problematic, since I wasn't using split(1) inside train_step I didn't notice this on an example I tried 😅

[ayaka14732]

from itertools import accumulate, chain, repeat
from jax.random import PRNGKey, split
from operator import itemgetter

seed = 42
keys = map(itemgetter(0), accumulate(chain((split(PRNGKey(seed)),), repeat(None)), lambda acc, _: split(acc[1])))

print(next(keys))  # [2465931498 3679230171]
print(next(keys))  # [1224796891 3487907634]
print(next(keys))  # [252818381 516428635]
print(next(keys))  # [3126834917  176854060]
print(next(keys))  # [ 207868211 2198494137]
print(next(keys))  # [358328843 191880770]
print(next(keys))  # [2134635248 4286345057]
print(next(keys))  # [2517085557 1569996244]

[jakevdp]

Be careful about solutions that rely on implicit side-effects, they may not work as expected under jit:

@jit
def next_key_jitted():
  return next(keys)

print(next_key_jitted())
# [2465931498 3679230171]
print(next_key_jitted())
# [2465931498 3679230171]

Here the side-effect is updating the global iterator state. For more on this see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions