Keys through vmap and scan

Split first, map second. Split a key into eight, vmap a sampler over the resulting key array, and the result is elementwise equal to the Python loop that draws from each key in turn. Not statistically indistinguishable. Equal.

That exactness belongs to the default generator, not to vmap. The jax.random reference states it as a property row in its implementation table, exact jax.vmap over keys, ticked for threefry with and without the partitionable flag and blank for both rbg implementations.

The counterexample runs here. Build eight rbg keys, vmap a scalar normal over them, and all eight values come from the first key alone: the vmapped result is elementwise equal to jax.random.normal(keys[0], (8,)), exactly as the docstring warns. Nothing raises, the shape is right, and the seven other keys were never consulted. If you ever switch implementation for TPU speed, this is the line whose meaning changes.

# Additionally, both ``"rbg"`` and ``"unsafe_rbg"`` behave unusually
# under ``jax.vmap``. When vmapping a random function over a batch
# of keys, its output values can differ from its true map over the
# same keys. Instead, under ``vmap``, the entire batch of output
# random numbers is generated from only the first key in the input
# key batch. For example, if ``keys`` is a vector of 8 keys, then
# ``jax.vmap(jax.random.normal)(keys)`` equals
# ``jax.random.normal(keys[0], shape=(8,))``.

import jax
import jax.numpy as jnp

keys = jax.random.split(jax.random.key(0), 8)
batched = jax.vmap(lambda k: jax.random.normal(k, (3,)))(keys)
looped = jnp.stack([jax.random.normal(k, (3,)) for k in keys])
print(batched.shape, bool(jnp.all(batched == looped)))

rbg = jax.random.split(jax.random.key(0, impl='rbg'), 8)
print(jax.vmap(jax.random.normal)(rbg))
print(jax.random.normal(rbg[0], (8,)))

# (8, 3) True
# [-1.7768954  -1.3365983   0.21346289 -0.16086003 -0.16176912 -0.2579535
#   0.8256621   0.03783947]
# [-1.7768954  -1.3365983   0.21346289 -0.16086003 -0.16176912 -0.2579535
#   0.8256621   0.03783947]

Hand vmap a single key with in_axes=None and it does precisely what you asked: broadcasts one key to every row. All eight rows then draw the same three numbers, the output shape is the shape you expected, and no rule was broken anywhere.

The museum's vmap-axes exhibit is the loud version of an in_axes mistake, where two arguments disagree about the batch size and vmap says so in the error. This one has no disagreement to report. A key is a shape-() value, so None is a legal thing to say about it.

Passing the split key array instead, with the default in_axes of 0, is what you meant, and the equality assertion from the previous lesson is how you keep it that way: no two rows elementwise equal.

import jax
import jax.numpy as jnp

key = jax.random.key(0)
rows = jax.vmap(lambda k, x: jax.random.normal(k, (3,)) + x, in_axes=(None, 0))(key, jnp.zeros((8, 3)))
print(rows[:2])
print(bool(jnp.all(rows == rows[0])))

keys = jax.random.split(key, 8)
ok = jax.vmap(lambda k, x: jax.random.normal(k, (3,)) + x)(keys, jnp.zeros((8, 3)))
print(bool(jnp.all(ok == ok[0])))

# [[ 1.8160863  -0.48262316  0.33988908]
#  [ 1.8160863  -0.48262316  0.33988908]]
# True
# False

Chapter 8 names both scan disciplines and calls them equally correct: carry the key and split it every step, or fold the step index into a base key that never moves. What the chapter does not say is that the two feed different counters to the same hash, so a loop that switches from one to the other changes every number it draws.

The first lesson printed both counters, and the loop below runs them six times each. The carried form draws from the second child of split, whose two words that same lesson printed, and every later step splits whatever the previous step handed on. The folded form hashes the pair (0, i) against a key that never changes, so step i is fixed by i alone. On this version the two streams share nothing: twelve distinct values on each side and a shared count of zero.

The carried form works only because a key is a legal scan carry. scan requires the carry to come back with the shape and dtype it went in with, and a key satisfies that: after six steps the carry still prints as dtype key<fry>, shape (), while the two words inside it have moved from [0 0] to [3110407274 4280739360]. Which discipline you want is decided by what happens when a run resumes partway through, and the next section measures that.

import jax
import jax.numpy as jnp

base = jax.random.key(0)

def carried(key, steps):
    def step(k, _):
        k, sub = jax.random.split(k)
        return k, jax.random.normal(sub, (2,))
    return jax.lax.scan(step, key, None, length=steps)

def folded(key, start, steps):
    def step(_, i):
        return None, jax.random.normal(jax.random.fold_in(key, i), (2,))
    return jax.lax.scan(step, None, jnp.arange(start, start + steps))

end_key, noise_a = carried(base, 6)
_, noise_b = folded(base, 0, 6)
print(noise_a)
print(noise_b)
print(bool(jnp.all(noise_a[0] == jax.random.normal(jax.random.split(base)[1], (2,)))))
print(len(set(map(float, noise_a.reshape(-1))) & set(map(float, noise_b.reshape(-1)))))
print(end_key.dtype, end_key.shape, jax.random.key_data(end_key))

# [[ 0.19307722 -0.52678293]
#  [ 0.00870701 -0.04888523]
#  [-0.89105326 -0.66184473]
#  [ 1.2091267   2.3117282 ]
#  [ 0.91798526  0.48332107]
#  [-0.21526915 -0.41612625]]
# [[-0.48121497 -0.01837499]
#  [ 1.4321449   1.3629805 ]
#  [ 1.1727159  -1.8279221 ]
#  [-0.05104028  0.68906456]
#  [ 0.05196318 -0.67536026]
#  [ 0.8082805  -0.21244425]]
# True
# 0
# key<fry> () [3110407274 4280739360]

Run six steps, then try to restart at step three, and the two disciplines come apart. The folded form needs two things: the base key and the integer 3. The three noise vectors that come back are elementwise equal to steps three through five of the uninterrupted run.

The carried form needs the key as it stood after step two, which means the key is part of what a checkpoint has to contain. Hand the function the original base key instead, and step three of the resumed run is elementwise equal to step zero of the first run. The loop restarted its stream, and the numbers look exactly as random as they did before.

Chapter twelve makes the rule out of this: params, optimizer state, step and key travel together. The run below is the test that fails when the key is left behind, and it is worth writing once for your own loop, because it is three assertions and it catches a class of bug that no loss curve will show you.

A resumed run that samples fresh noise from step zero is not a resumed run.

import jax
import jax.numpy as jnp

base = jax.random.key(0)

def carried(key, steps):
    def step(k, _):
        k, sub = jax.random.split(k)
        return k, jax.random.normal(sub, (2,))
    return jax.lax.scan(step, key, None, length=steps)

def folded(key, start, steps):
    def step(_, i):
        return None, jax.random.normal(jax.random.fold_in(key, i), (2,))
    return jax.lax.scan(step, None, jnp.arange(start, start + steps))

_, whole_a = carried(base, 6)
_, whole_b = folded(base, 0, 6)

key_at_3, _ = carried(base, 3)
_, from_base = carried(base, 3)
_, from_saved = carried(key_at_3, 3)
_, resumed_b = folded(base, 3, 3)

print(bool(jnp.all(from_base[0] == whole_a[0])))
print(bool(jnp.all(from_saved == whole_a[3:])))
print(bool(jnp.all(resumed_b == whole_b[3:])))

# True
# True
# True

Two folds compose. Fold the example id into the base key, fold the step into the result, and every pair of coordinates gets its own key with no carry to thread and no split width to decide in advance. Twelve draws across a three by four grid come back as twelve distinct values.

The reason this composes is a staticness rule you can see in the error. split takes a count that turns into a shape, so under jit it demands a concrete Python integer and says Shapes must be 1D sequences of concrete values of integer type when it does not get one. fold_in only hashes its argument, so a traced integer is fine. Anything indexed by a runtime quantity, a step, an example id, a layer number, is fold_in work.

Coordinates you can name are coordinates you can resume from, which is the property the previous section measured. Deriving a key from where a piece of work sits, rather than from how many draws happened before it, is what makes a sampler restartable at any point.

import jax
import jax.numpy as jnp

key = jax.random.key(0)

def stream(example, step):
    return jax.random.fold_in(jax.random.fold_in(key, example), step)

grid = jax.vmap(lambda e: jax.vmap(lambda s: jax.random.normal(stream(e, s), ()))(jnp.arange(4)))(jnp.arange(3))
print(grid)
print(len(set(map(float, grid.reshape(-1)))))

try:
    jax.jit(lambda n: jax.random.split(key, n))(3)
except TypeError as err:
    print(str(err).splitlines()[0])

# [[-0.9984514   0.4432096  -0.2424067  -1.1119174 ]
#  [-0.04697571  1.2683467   1.816136   -0.6499298 ]
#  [ 1.3010085   0.18747099  1.159638    0.10436483]]
# 12
# Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace>,).