Common Mistakes and Reproducibility

PRNG keys are simple, but there are a few mistakes that even experienced JAX users make. Let's walk through them so you can avoid them.

Mistake 1: Reusing Keys (Correlated Randomness)

import jax
import jax.numpy as jnp

key = jax.random.key(0)

# WRONG: same key → same values → correlated randomness
w1 = jax.random.normal(key, (100, 50))
w2 = jax.random.normal(key, (50, 10))
# w1 and w2 share the same random pattern (scaled to shape)!

# RIGHT: split first
k1, k2 = jax.random.split(key)
w1 = jax.random.normal(k1, (100, 50))
w2 = jax.random.normal(k2, (50, 10))
# Fully independent

Mistake 2: Forgetting to Split in Loops

key = jax.random.key(0)

# WRONG: same key every iteration
samples_bad = []
for i in range(5):
    samples_bad.append(jax.random.normal(key, (3,)))
# All five samples are IDENTICAL

# RIGHT: split at each iteration
samples_good = []
for i in range(5):
    key, subkey = jax.random.split(key)
    samples_good.append(jax.random.normal(subkey, (3,)))
# All five samples are different

Mistake 3: Using Python's random or NumPy random with JAX

import random
import numpy as np

# WRONG: these are not JIT-compatible
# x = random.random()           # Python random — not traced
# x = np.random.randn(3)        # NumPy random — not traced

# RIGHT: always use jax.random inside jitted functions
@jax.jit
def generate(key):
    return jax.random.normal(key, (3,))

Reproducibility Guarantees

JAX provides strong reproducibility: same key + same function = same output, every time, on the same hardware and JAX version.

key = jax.random.key(123)

# Run 1
a = jax.random.normal(key, (5,))

# Run 2 — same key, same result
b = jax.random.normal(key, (5,))

assert jnp.array_equal(a, b)  # True — always

# This works even under jit
@jax.jit
def sample(key):
    return jax.random.normal(key, (1000,))

c = sample(key)
d = sample(key)
assert jnp.array_equal(c, d)  # True

Here's a complete example tying it all together — implementing dropout and data augmentation with proper PRNG management:

import jax
import jax.numpy as jnp

def create_train_state(key, input_dim, hidden_dim, output_dim):
    """Initialize model with reproducible randomness."""
    k1, k2, k3 = jax.random.split(key, 3)
    params = {
        'w1': jax.random.normal(k1, (input_dim, hidden_dim)) * 0.01,
        'b1': jnp.zeros(hidden_dim),
        'w2': jax.random.normal(k2, (hidden_dim, output_dim)) * 0.01,
        'b2': jnp.zeros(output_dim),
    }
    return params

@jax.jit
def train_step(params, batch, dropout_key):
    """Single training step with dropout."""
    x, y = batch
    k1, k2 = jax.random.split(dropout_key)

    def loss_fn(params):
        h = jax.nn.relu(x @ params['w1'] + params['b1'])
        mask = jax.random.bernoulli(k1, 0.7, h.shape)
        h = jnp.where(mask, h / 0.7, 0.0)
        logits = h @ params['w2'] + params['b2']
        return jnp.mean((logits - y) ** 2)

    loss, grads = jax.value_and_grad(loss_fn)(params)
    params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
    return params, loss

# Training loop with proper key management
key = jax.random.key(0)
key, init_key = jax.random.split(key)
params = create_train_state(init_key, 784, 256, 10)

for step in range(1000):
    key, step_key = jax.random.split(key)
    # params, loss = train_step(params, batch, step_key)