Pseudorandom numbers

Pseudorandom numbers#

In most numerical libraries, drawing a random number means consulting a hidden generator that updates itself between calls. JAX has no such generator, and no hidden update. Randomness in JAX is driven by keys: ordinary immutable array values that you pass around explicitly. Every function in jax.random is a pure, deterministic function of its key argument: call it twice with the same key and you get the same number twice, by design.

Why not a global generator?#

NumPy’s numpy.random is the familiar stateful design: a global generator, seeded once, silently advanced by every sampling call:

import numpy as np

np.random.seed(0)
print(np.random.uniform())
print(np.random.uniform())  # different: the hidden state advanced

0.5488135039273248
0.7151893663724195

The problem is that “the hidden state advanced” makes your results depend on the exact order and number of sampling calls everywhere in your program:

np.random.seed(0)

def bar(): return np.random.uniform()
def baz(): return np.random.uniform()

def foo(): return bar() + 2 * baz()

print(foo())

1.9791922366721637

This value is only reproducible because NumPy promises to run bar() before baz(). Such sequencing promises are what JAX needs to avoid: an optimizing compiler should be free to reorder and parallelize work across devices. JAX needs random number generation that is reproducible, parallelizable, and vectorizable, which rules out sampling functions that secretly read and write shared state.

The solution isn’t to make the generator state an explicit argument that you shuttle in and out of functions. Instead, sampling is just a pure function of key values, and no updated state is threaded back out.

Keys are values#

Create a key from an integer seed with jax.random.key():

import jax
from jax import random

key = random.key(42)
key

Array((), dtype=key<fry>) overlaying:
[ 0 42]

A key is a rank-0 array (shape ()) with a special element type: its JAX type is key<fry>[], where key<fry> names the default Threefry PRNG implementation.

print(jax.typeof(key))

key<fry>[]

Passing a key to a sampling function doesn’t modify it or “use it up” in any physical sense; the sample is a deterministic function of the key:

print(random.normal(key))
print(random.normal(key))  # same key, same value — necessarily

-0.028304616
-0.028304616

This means reproducibility is automatic: results depend only on the key values your program constructs, never on execution order, call counts, or which device ran what.

The flip side is that distinct random numbers require distinct keys, which leads to the one rule of JAX randomness:

Never reuse a key (unless you want identical outputs). Feeding the same key to two different samplers produces correlated results, depriving your program of lifegiving chaos.

Deriving new keys#

To get fresh keys, derive them from a key you already have. jax.random.split() deterministically produces any number of new keys, each of which can be used to generate statistically independent samples:

key = random.key(42)
key, subkey = random.split(key)
print(random.normal(subkey))

0.60576403

split produces as many keys as you ask for in one shot:

key = random.key(42)
subkeys = random.split(key, num=4)
[float(random.normal(k)) for k in subkeys]

[0.07592553645372391,
 0.6057640314102173,
 0.4323064982891083,
 -0.281894713640213]

And jax.random.fold_in() derives a new key from a key and an integer, which is ideal for generating a per-step or per-example key without carrying any key-threading through your loop:

key = random.key(42)
for step in range(3):
  step_key = random.fold_in(key, step)
  print(f"step {step}: {random.normal(step_key)}")

step 0: 0.07592553645372391
step 1: 0.6057640314102173
step 2: 0.4323064982891083

Note the shape of this pattern: every step_key is derived directly from one parent key, not from its predecessor. The next section explains why.

Keep the key tree wide, not deep#

Your program’s keys form a tree, rooted at the seed, growing by either split or fold_in operations. A bad pattern is to grow that tree as a long chain, splitting off each step’s key from the previous step’s:

for step in range(num_steps):
  key, subkey = random.split(key)   # each key derived from the last: avoid!
  ...

Prefer wide trees instead: all step keys hanging off one parent, via a single split(key, num_steps) or via fold_in(key, step) as above. The chained version has two problems, one computational and one statistical:

  • It serializes. Each key depends on the previous one, so a chain of a million steps means a million sequential hash applications. Wide key derivation is a single batched operation, free to vectorize and parallelize.

  • It courts collisions. For a fixed key, the PRNG’s underlying hash is a pseudorandom permutation of its input, so the keys produced by one split (or by fold_in over distinct integers) are guaranteed distinct. But viewed as a function of the key, the hash is not a permutation: it behaves like a random function. Every derivation hop is therefore an independent chance for two keys in your tree to coincide, and over a long chain in the default 64-bit key space, collision probability accumulates toward the birthday bound. A collision means identical random streams from the point of collision onward.

A handful of chained splits is harmless; the collision math only bites at scale, and plenty of correct code splits a key a few times in sequence. But for anything proportional to the length of training or the size of a dataset, derive keys widely from a common parent.

No sequential equivalence#

NumPy guarantees that sampling N numbers one at a time yields the same sequence as sampling N at once. JAX deliberately makes no such promise:

key = random.key(42)
subkeys = random.split(key, 3)
print("individually:", np.stack([random.normal(k) for k in subkeys]))

key = random.key(42)
print("all at once: ", random.normal(key, shape=(3,)))

individually: [0.07592554 0.60576403 0.4323065 ]
all at once:  [-0.02830462  0.46713185  0.29570296]

Sequential equivalence would impose the kind of ordering constraint JAX’s design exists to avoid. Giving it up means samples drawn from independent keys don’t depend on each other in any order, so generation can be freely vectorized and sharded.

Since keys are just arrays, they compose with everything else in JAX. You can vmap a sampler over a batch of keys:

import jax
jax.vmap(random.normal)(subkeys)

Array([0.07592554, 0.60576403, 0.4323065 ], dtype=float32)

With the default PRNG implementation, this is exactly equivalent to calling random.normal on each key separately; vectorizing over keys doesn’t change the values.

Note

These docs use the typed keys created by jax.random.key(). You may also encounter older code using jax.random.PRNGKey(), which produces a raw uint32 array — it still works, but it’s easy to misuse (nothing stops you from doing arithmetic on it) and it doesn’t record which PRNG implementation it belongs to. Prefer jax.random.key in new code, and convert at boundaries with jax.random.key_data() and jax.random.wrap_key_data() when interfacing with systems that need raw arrays. See the typed PRNG keys JEP for the full story.

Design and implementations#

JAX’s PRNG is a counter-based Threefry hash combined with a functional splitting model, chosen so that generation has no sequencing constraints at all. See the PRNG design JEP for the design rationale.

Threefry is the default of several available implementations. Alternatives (selected per-key via the impl argument to jax.random.key(), or globally via the jax_default_prng_impl config) trade off generation speed on TPUs, shardability and sharding-invariance, bit-for-bit identical results across platforms, and exact vmap-over-keys semantics. The jax.random module documentation has the full comparison table; the default is the right choice unless PRNG generation shows up in your profiles.

jax.random itself offers samplers for a wide range of distributions — uniform, normal, categorical, permutations, and many more — all taking a key as their first argument.

Next steps#

Keys handle randomness while keeping every function pure. The remaining topic is state: values that evolve as a program runs, and genuine in-place mutation. It’s covered in Stateful computations.