Randomness

Call numpy.random.normal() twice and you get two different answers, because NumPy keeps a global seed mutating quietly in the background. JAX has no such thing. Every sampling function takes an explicit key, and the same key with the same shape produces identical bits, every call, every run, every machine, for a given backend. Reproducibility isn't a flag you remember to set. It's just how the system works.

The same key gives the same bits, forever.

The reason traces back to chapter one's invariant: values, not places. A global RNG dial is hidden mutable state, exactly the kind of thing that makes a traced function lie about what it does. It would also make vmap and jit order-dependent: run the same batched computation twice and get different random draws depending on how the compiler chose to schedule things. An explicit key sidesteps all of that, because the randomness lives in a value you pass around like any other array.

jax.random.key(0) builds one of these keys, a small, typed object, the modern form you should reach for. Older code still shows jax.random.PRNGKey(...), which returns a plain uint32 array instead of a typed key; both work, but the typed key is what current JAX expects, and what this chapter uses throughout.

A key follows one rule, and it's worth memorizing before you write a single sampling call: use it exactly once. Either a sampler consumes it directly, or you split it into children and use those instead. Reuse a key for two different draws and the draws aren't independent anymore. They're correlated, silently, and nothing in the code will warn you.

jax.random.split(key, n) is how you get n independent keys out of one. fold_in(key, step) does something related but different: it derives a new key deterministically from an existing one plus an integer, like a step number or a layer id, without you having to carry a growing chain of split keys through the rest of the function. Both are correct ways to get more randomness out of one seed; which one reads better depends on the shape of the code around it.

import jax

key = jax.random.key(0)
k1, k2 = jax.random.split(key)
print(jax.random.normal(k1, (2,)))   # independent of k2's draw
print(jax.random.normal(k2, (2,)))
print(jax.random.normal(k1, (2,)))   # identical to the first line, forever

Batched sampling follows the same discipline, just wired through vmap. Split a key into however many independent draws you need, then map a sampler over the resulting array of keys. The result is one fused program producing N independent samples, not N sequential calls to a sampler.

Inside a scan, you have two options for the same problem. Carry the key in the loop's carry and split it fresh on every step, or fold_in the step index each time and skip the carry entirely. Both give you correct, independent randomness per step. fold_in tends to read better for loops you expect to resume partway through, chapter twelve leans on exactly this, because the key for step k doesn't depend on having run steps 0 through k-1 first.

import jax

key = jax.random.key(0)
keys = jax.random.split(key, 32)
samples = jax.vmap(lambda k: jax.random.normal(k, (3,)))(keys)
print(samples.shape)   # (32, 3): 32 independent draws, one fused program