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
JAX has no global seed. All random operations require an explicit PRNG Key array (jax.Array).
Calling a random sampler with the same key produces identical numbers. Always use jax.random.split to derive new subkeys.
PRNG keys are standard array payloads; they work seamlessly inside jax.jit, jax.vmap, and jax.grad.
Use the key, subkey = jax.random.split(key) pattern in loops to keep key streams organized.
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.
