JAX provides a comprehensive set of random distributions in jax.random. Every function takes a PRNG key as its first argument — no exceptions.
import jax
import jax.numpy as jnp
key = jax.random.key(42)
keys = jax.random.split(key, 8)
# Uniform [0, 1)
uniform = jax.random.uniform(keys[0], shape=(3,))
# [0.299, 0.784, 0.033]
# Normal (mean=0, std=1)
normal = jax.random.normal(keys[1], shape=(3,))
# [-0.272, 1.085, -0.533]
# Bernoulli (coin flips)
coins = jax.random.bernoulli(keys[2], p=0.7, shape=(5,))
# [True, True, False, True, True]
# Categorical (sample from discrete distribution)
logits = jnp.array([1.0, 2.0, 0.5])
category = jax.random.categorical(keys[3], logits, shape=(4,))
# [1, 1, 0, 1] — category 1 most likely
# Truncated normal (clipped to [lower, upper])
trunc = jax.random.truncated_normal(keys[4], lower=-2.0, upper=2.0, shape=(3,))
# Randint (random integers)
ints = jax.random.randint(keys[5], shape=(3,), minval=0, maxval=10)
# Permutation (shuffling)
shuffled = jax.random.permutation(keys[6], jnp.arange(5))
# Exponential
exp_samples = jax.random.exponential(keys[7], shape=(3,))Here's a side-by-side comparison with NumPy equivalents:
# NumPy # JAX
# np.random.randn(3, 4) jax.random.normal(key, (3, 4))
# np.random.rand(3, 4) jax.random.uniform(key, (3, 4))
# np.random.randint(0, 10, (3,)) jax.random.randint(key, (3,), 0, 10)
# np.random.choice(arr, size=5) jax.random.choice(key, arr, (5,))
# np.random.shuffle(arr) jax.random.permutation(key, arr)
# Key difference: JAX always needs a key, and never mutates input💡 Why This Matters
In PyTorch, you'd use torch.randn(3, 4) which also uses global state (but lets you set torch.manual_seed()). JAX's explicit approach is more verbose, but it makes random operations deterministic under all transformations — including jit, vmap, pmap, and grad. If you're doing distributed training across 8 GPUs, explicit keys ensure each device gets truly independent randomness without subtle correlation bugs.
For neural network initialization, JAX provides convenient initializers in libraries like Flax, but they all use PRNG keys under the hood:
# Xavier/Glorot initialization from scratch
def glorot_normal(key, shape):
fan_in, fan_out = shape[-2], shape[-1]
std = jnp.sqrt(2.0 / (fan_in + fan_out))
return jax.random.normal(key, shape) * std
key = jax.random.key(0)
w = glorot_normal(key, (256, 128))
print(f"Mean: {w.mean():.4f}, Std: {w.std():.4f}")
# Mean: ~0.0, Std: ~0.072