If you've used NumPy, you're used to calling np.random.randn(3) and getting random numbers. Easy. But there's a hidden catch: NumPy maintains a global random state — a mutable counter that changes every time you draw a random number. That hidden state is an invisible side effect, and side effects are poison for JAX's functional programming model.
Think about it this way: a pure function always returns the same output for the same input. But if a function secretly reads from and writes to a global counter, the same function call can produce different results each time. That makes it fundamentally incompatible with jit, vmap, and pmap — transformations that need to reason about what a function does by looking only at its inputs and outputs.
# NumPy: hidden global state — NOT functional
import numpy as np
np.random.seed(42)
print(np.random.randn()) # 0.4967...
print(np.random.randn()) # -0.1383... (different! state mutated)
# JAX: explicit state — fully functional
import jax
import jax.numpy as jnp
key = jax.random.key(42)
print(jax.random.normal(key)) # always the same value
print(jax.random.normal(key)) # exact same value again!In JAX, random state is an explicit value you pass around — called a PRNG key. The same key always produces the same random numbers. To get different numbers, you create new keys by splitting the current one. There's no hidden mutation anywhere.
💡 Why This Matters
This isn't just philosophical purity. Without explicit random state, jit compilation could cache the first random result and return it forever. vmap would give every element of a batch the same "random" values. pmap across devices would produce correlated randomness. Explicit PRNG keys make all of these work correctly and reproducibly.
⚠️ Pure Function Check
A function that calls np.random.randn() is not pure — it reads and mutates global state. A function that takes a PRNG key and calls jax.random.normal(key) is pure — same key in, same numbers out, no side effects.
The mental shift is small but important: in NumPy, randomness is something that happens to you. In JAX, randomness is something you explicitly control. Once you internalize this, JAX's random API feels natural and gives you stronger reproducibility guarantees than any other framework.