Key Management Patterns

Good key management is the difference between correct and subtly broken randomness. Here are the patterns every JAX programmer needs to know.

Pattern 1: Split Before Use

Always split a key before consuming it. Never reuse a key for two different random operations.

import jax
import jax.numpy as jnp

def init_layer(key, in_dim, out_dim):
    """Initialize a single layer with split keys."""
    k1, k2 = jax.random.split(key)
    weights = jax.random.normal(k1, (in_dim, out_dim)) * 0.01
    biases = jax.random.normal(k2, (out_dim,)) * 0.01
    return {'w': weights, 'b': biases}

def init_network(key, layer_sizes):
    """Initialize a full network, splitting keys for each layer."""
    params = []
    for i in range(len(layer_sizes) - 1):
        key, subkey = jax.random.split(key)
        params.append(init_layer(subkey, layer_sizes[i], layer_sizes[i+1]))
    return params

key = jax.random.key(42)
params = init_network(key, [784, 256, 128, 10])
print(f"Layer 0 weights shape: {params[0]['w'].shape}")  # (784, 256)

Pattern 2: Pass Keys Explicitly

If a function needs randomness, it should take a key as an argument. Never store keys in global variables.

# WRONG: global key — leads to reuse bugs
# global_key = jax.random.key(0)  # Don't do this!

# RIGHT: pass key through the call chain
def dropout(x, key, rate=0.5):
    """Apply dropout with an explicit key."""
    mask = jax.random.bernoulli(key, 1.0 - rate, x.shape)
    return jnp.where(mask, x / (1.0 - rate), 0.0)

def forward(params, x, key):
    """Forward pass with dropout — key passed explicitly."""
    k1, k2 = jax.random.split(key)
    h = jnp.tanh(x @ params[0]['w'] + params[0]['b'])
    h = dropout(h, k1, rate=0.3)
    h = jnp.tanh(h @ params[1]['w'] + params[1]['b'])
    h = dropout(h, k2, rate=0.3)
    return h @ params[2]['w'] + params[2]['b']

Pattern 3: Different Key Each Training Step

For stochastic operations like dropout, you need a new key at every step. The cleanest approach uses jax.random.fold_in:

# Option A: fold_in the step number (deterministic, no key threading)
base_key = jax.random.key(0)
for step in range(1000):
    step_key = jax.random.fold_in(base_key, step)
    # step_key is unique for each step, reproducible from base_key + step

# Option B: split at each step (standard threading)
key = jax.random.key(0)
for step in range(1000):
    key, subkey = jax.random.split(key)
    # use subkey for this step

Think of keys like spending money: once you use a key, it's "spent." Split to get new keys, and never spend the same one twice.