Managing PRNG Keys in JAX

In the previous articles, we have explored the fundamental principles of JAX from functional purity to transformations to how it handles arrays to PyTrees. However, in real-world ML, when working with JAX especially from initializing model weights, applying dropout, shuffling datasets, among others, you'll notice a JAX unique design pattern using Pseudo-Random Number Generation (PRNG)

Those with a background in PyTorch or NumPy, JAX approach to randomness will feel surprisingly strict as there's no global seed or random state. In this article, we shall explore JAX's PRNG system from why global random states break functional purity to understanding key splitting patterns in development. Then later, I'll try to illustrate how to manage them production pipelines and stateful models.

Why JAX rejects global random states

In standard numpy or PyTorch, generating random numbers relies on a global stateful random number generator.

Consider the NumPy code below

import numpy as np

np.random.seed(42)
print(np.random.normal()) # Generates a random number and mutates the global state
print(np.random.normal()) # Generates a different number because a different global state

Output

0.4967141530112327
-0.13826430117118466

Every time you call np.random.normal(), it returns a pseudo-random value and rotates the internal global state counter so the next call produces a different figure. But as established in the first study JAM, JAX functions must remain pure i.e. given the same inputs, the same outputs are expected without any side effects on the output. A hidden mutating global state completely violates the principle of functional purity.

Furthermore, global state makes parallel execution nondeterministic. If two parallel threads request a random number from a single global state, the execution order determines which thread gets which number.

To guarantee reproducibility, parallel safety, and functional purity, JAX makes random states explicit using PRNG Keys.

Abstractions with jax.random.key and jax.random.split

In JAX, randomness is functional. A random function takes a mathematical key array as an explicit argument, deterministically transforms that key, and returns the requested random numbers.

If you pass the exact same key to a random function twice, you will get the same output twice, as demonstrated in the example below

import jax
import jax.numpy as jnp

#create an initial key
key = jax.random.key(42)

#generating randoms with the same key yields the same output
side_a = jax.random.normal(key, (3,))
side_b = jax.random.normal(key, (3,))

print("Draw 1:", side_a)
print("Draw 2:", side_b)
print("Are they identical?", jnp.array_equal(side_a, side_b))

Output

Draw 1: [-0.02830462  0.46713185  0.29570296]
Draw 2: [-0.02830462  0.46713185  0.29570296]
Are they identical? True

Splitting Keys

To generate new, independent random numbers, you must fold or split your existing key into new unique keys using jax.random.split.

import jax

#initialize the root key
root_key = jax.random.key(420)

# split into a usable subkey and hold out the root_key
root_key, subkey = jax.random.split(root_key)

#use the subkey for generating any randoms
new_random = jax.random.normal(subkey, (3,))
print(f"Draw with subkey: {new_random}")

Output

Draw with subkey: [-1.4182527   0.12200668  0.8509691 ]

As a rule of thumb

Never reuse a key. Once a key generates random numbers or splits into subkeys, discard it and use a fresh subkey for subsequent operations.

Managing Keys in Development

When prototyping algorithms or training scripts locally, managing key splitting manually is simple once you establish disciplined coding habits.

Splitting and consuming keys in loops

In sequential scripts or training loops, keep a single key variable in scope. At each step, split the current key into a fresh step_key to consume and update key with the remaining stream.

import jax

root_key = jax.random.key(420)

# training steps with random operations
for step in range(3):
    root_key, step_key = jax.random.split(root_key)

    # pass the key into your model or data transformation
    random_noise = jax.random.normal(step, (2,))
    print(f"Step {step + 1} noise sample: {random_noise}")

Output

Step 1 noise sample: [-1.4182527   0.12200668]
Step 2 noise sample: [-0.06860777  0.04123694]
Step 3 noise sample: [-1.8361075  1.5372876]

Multi-Subkey Splitting

If a single function or initialization step requires multiple random operations, you can split a key into ‭N‬ subkeys simultaneously. Consider the example below

import jax

root_key = jax.random.key(256)

#Generate distinct subkeys at once
subkeys = jax.random.split(root_key, num=4)

print("Subkeys array shape:", subkeys.shape)

Output

Subkeys array shape: (4,)

Vectorizing and Compiling Randomness

Because PRNG keys are standard jax.Array objects under the hood, they fit directly into JAX's core transformations i.e. jit and vmap

Vectorizing Random Sampling with jax.vmap

Suppose you want to generate independent random trajectories or samples across a batch. Instead of running a Python loop, you can split a key into a batch of keys and use jax.vmap

import jax
import jax.numpy as jnp

root_key = jax.random.key(254)

# Get 5 keys
batch_keys = jax.random.split(root_key, num=5)

# Vectorize jax.random.normal across axis 0 of batch_keys
vmapped_normal = jax.vmap(lambda k: jax.random.normal(k, shape=(3, )))

# Get 5 distinct samples of shape (3,) in parallel
batch_samples = vmapped_normal(batch_keys)
print("Batch random samples shape:", batch_samples.shape)

Output

Batch random samples shape: (5, 3)

Randomness inside jax.jit

Passing keys into jax.jit compiled functions works transparently. Since keys are pure arrays, compiling functions that accept keys introduces zero side effects.

import jax
import jax.numpy as jnp

@jax.jit
def apply_dropout(x, key, dropout_rate=0.2):
    keep_prob = 1.0 - dropout_rate

    # Generate binary mask
    mask = jax.random.bernoulli(key, p=keep_prob, shape=x.shape)

    #Scale remaining values to preserve expected value
    return jnp.where(mask, x / keep_prob, 0)

key = jax.random.key(99)
data = jnp.ones((4, 4))

key, subkey = jax.random.split(key)
dropped_data = apply_dropout(data, subkey)
print("Dropped Data:\n", dropped_data)

Output

Dropped Data:
 [[1.25 1.25 1.25 0.  ]
 [1.25 1.25 1.25 1.25]
 [1.25 1.25 1.25 0.  ]
 [1.25 0.   0.   1.25]]

PRNG management when things get serious

While manually splitting with key, subkey = jax.random.split(key) works well for simple scripts, doing this manually across involved neural network architectures becomes unwieldy and error-prone.

Consider the following scenarios from the codebases below

Functional key passing

In pure JAX systems, training steps accept a top-level key as an argument alongside model parameters and batch data. The step function handles the splitting internally and passes dedicated keys down to sub-routines as exhibited in the code sample below.

import jax
import jax.numpy as jnp

def model_forward(params, x, key, is_training=True):
    if not is_training:
        return jnp.dot(x, params['w'])

    w_key, dropout_key = jax.random.split(key)
    out = jnp.dot(x, params['w'])

    # Pass dropout_key to internal dropout logic
    mask = jax.random.bernoulli(dropout_key, p=0.8, shape=out.shape)
    
    return jnp.where(mask, out / 0.8, 0.0)

@jax.jit
def train_step(params, opt_state, x_batch, y_batch, key):
    # Split key for step operations
    step_key, dropout_key = jax.random.split(key)

    # Compute predictions using explicit key
    preds = model_forward(params, x_batch, dropout_key, is_training=True)
    loss = jnp.mean((preds - y_batch) ** 2)

    # ... Gradients and updates ...
    return params, opt_state, loss
 

Stateful PRNG Generators in Class Abstractions

Higher level libraries like Equinox or Flax encapsulate key managementwhen building layer stacks. For example Equinox provides a clean state management pattern but if you're feeling creative, a custom PRNG generator (as exhibited in the previous article) is also an option.

Consider the code sample below

import jax

class PRNGSequence: 
    """ for sequential key splitting """
    def __init__(self, seed_or_key):
        if isinstance(seed_or_key, int):
            self._key = jax.random.key(seed_or_key)
        else: 
            self._key = seed_or_key

    def __call__(self):
        """ returns a fresh subkey while maintaining the internal key state """
        self._key, subkey = jax.random.split(self._key)
        return subkey

# Driver usage
rng = PRNGSequence(seed=420)

# Each call returns a new key
key_0 = rng()
key_1 = rng()
key_2 = rng()

print(f"Key 0 {key_0} - Key 1 {key_1} - Key 2 {key_2}")

Output

Key 0 Array((), dtype=key<fry>) overlaying:
[1499546378 2761259651] - Key 1 Array((), dtype=key<fry>) overlaying:
[3856841681 3513101443] - Key 2 Array((), dtype=key<fry>) overlaying:
[2301777456 4258758523]

By decoupling the key splitting code from model forward passes, your architecture code remains legible while guaranteeing deterministic execution.

All in all

1. JAX has no global seed. All random operations require an explicit PRNG Key array (jax.Array).

2. Calling a random sampler with the same key produces identical numbers. Always use jax.random.split to derive new subkeys.

3. PRNG keys are standard array payloads; they work seamlessly inside jax.jit, jax.vmap, and jax.grad.

4. Use the key, subkey = jax.random.split(key) pattern in loops to keep key streams organized.

5. Pass root keys explicitly into top-level @jax.jit step functions, or wrap key generation in a disciplined abstraction like a PRNGSequence driver or framework modules (Equinox/Flax).

See you all in the next one.