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.

JAX’s solution isn’t simply to make the generator state an explicit argument that you shuttle in and out of functions. Instead, sampling is a pure function of a key, and no updated state comes 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 life-giving chaos.

Deriving new keys#

To get fresh keys, derive them from a key you already have. jax.random.split() deterministically derives 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.605764

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.6057639718055725,
 0.4323064684867859,
 -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 threading a key 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.6057639718055725
step 2: 0.4323064684867859

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.

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

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.605764   0.43230647]
all at once:  [-0.02830462  0.46713176  0.29570293]

Sequential equivalence would impose the kind of ordering constraint JAX’s design exists to avoid. Without it, samples drawn from different keys have no ordering relationship at all, 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:

jax.vmap(random.normal)(subkeys)

Array([0.07592554, 0.605764  , 0.43230647], 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.

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.