JAX Pathways
Core Foundation · Local CPULaptop CPULesson 1 of 4 in Phase 03

Phase 03 / State, randomness & control flow

Random keys without surprises

Revision v0.7.0 ·

Lesson · about 60 minutes · CPU practice

Split keys and show reproducibility without accidental key reuse.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: JAX transformations
Hardware: CPU

Download notebook
Other downloadsPython worked examplesLesson text
Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder
py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment
python phases/03-state/01-random-keys-without-surprises/code/main.py

PowerShell command · repository folder
.\.venv\Scripts\python.exe phases/03-state/01-random-keys-without-surprises/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Check: Replay equal: True; Fresh draw equal: False.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 03.01 Stage 1 of 5 · Understand
1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence
Stage 1 of 5 · Understand the MechanismCore mental model, equations & visual walkthrough

What you will be able to do

  • Treat a key as an explicit input
  • Carry the continuation key explicitly
  • Separate replay from a statistical claim
  • Allocate randomness by sample identity

The problem

You saved the seed, yet your simulation repeats exactly the same noise at every step. Let’s follow the random key. In JAX, using a key again repeats the same draw; splitting it gives you keys for separate draws. We’ll practice both deliberate replay and fresh sampling before adding randomness to a training loop.

The idea

A random sampler in JAX receives an explicit key. Reusing the same key with the same sampler repeats the same computation. Splitting gives us separately owned keys, so the program can show which operation receives which randomness.

A key is an explicit input
  1. Parent keykey

    Keep one owner for the sequence.

  2. Splitnext_key, sample_key

    Use one child to continue, another for this sample.

  3. Next callnext_key → split

    Pass the continuation key forward.

Reusing sample_key reproduces the same random computation. Splitting does not mutate the parent key; decide which child belongs to each computation.

Conceptual illustration. Verify the computation in the lesson’s Python experiments.

A key is an input, not a moving cursor

Imagine a training step that needs dropout randomness and a key for the next step. Split the incoming key, use one child for dropout, and return the other as the continuation key. The drawing should branch once; it should not show a sampler secretly advancing the parent key.

If two augmentation calls receive the same key and otherwise identical inputs, they can produce the same augmentation. Different variable names do not make the underlying key values different. This is a data-flow bug that ordinary shape checks will miss.

For recovery, save the continuation key at the same boundary as model and data position. Reconstructing an initial seed and restarting the random stream is not equivalent to restoring the next key after many training steps.

St+1=update⁡(St,Xt),pure transition (St,Xt)↦(St+1,Yt)S_{t+1} = \operatorname{update}(S_t, X_t), \qquad \text{pure transition } (S_t, X_t) \mapsto (S_{t+1}, Y_t)

Treat a key as an explicit input

A JAX random operation is deterministic given its key and other arguments. The key is not a mutable generator that automatically advances after a call. Calling normal twice with the same key, shape and dtype reproduces the draw. That property is useful for replay and dangerous when accidental reuse makes samples correlated.

A seed is only the initial description of an experiment. Splitting describes where subsequent random choices come from. Write down which function owns the continuation key and which subkey is consumed by a draw. Do not pass one subkey to two independent components.

seed → root key
root key --split--> continuation key + sample key
sample key --normal(shape=(3,))--> one sample

Carry the continuation key explicitly

In the worked program, key, sample_key = split(key) replaces the variable key with a continuation value. The old key was not mutated; Python rebound a name. sample_key belongs to the immediate draw. The next split uses the retained continuation, not sample_key.

This convention gives each training or simulation step a small state transition: receive key → split → sample → return next key. Later, checkpoints need to store that next key alongside model and optimizer state. Storing parameters alone cannot reproduce the next random batch.

Separate replay from a statistical claim

For this fixed seed, fresh subkeys produce visibly different draws. That observation is not proof of statistical independence, and different random draws can theoretically coincide. Equality of a single pair is therefore not a general test of randomness quality.

The reliable replay test is to reconstruct the same program from the same seed and compare its results. A change to key allocation can alter all later samples. Do not promise identical streams across arbitrary JAX versions, PRNG implementations, shapes or devices.

Allocate randomness by sample identity

Splitting a root into four keys is useful for one fixed batch. If the batch membership changes, assignment by position may associate a different draw with a particular sample. fold_in lets you combine a stable integer identity with a base key, which can make per-example randomness stable under reordering.

Keep identifiers unique where distinct randomness is intended. Reusing an identifier repeats its derived key. This is deterministic bookkeeping, not a security primitive or a substitute for recording the PRNG implementation.

Reused keys replay the same sample

Predict before looking: Will a repeated draw lie on top of the original?

Reused keys replay the same sample. The horizontal axis selects a coordinate within a three-element random draw; the vertical axis gives its sampled value. The solid first-draw line and dashed same-key replay line lie on top of each other. The dash-dot line comes from a new split key. The replay repeats approximately \((0.606,0.799,-0.909)\), while the split-key draw is approximately \((-0.211,-1.363,-0.045)\). Two curves are hard to distinguish precisely because their coordinates are equal.
Plot from a recorded CPU computation · JAX 0.9.2 · 1 CPU device

Read the figure

The horizontal axis selects a coordinate within a three-element random draw; the vertical axis gives its sampled value. The solid first-draw line and dashed same-key replay line lie on top of each other. The dash-dot line comes from a new split key.

The replay repeats approximately (0.606,0.799,−0.909)(0.606,0.799,-0.909), while the split-key draw is approximately (−0.211,−1.363,−0.045)(-0.211,-1.363,-0.045). Two curves are hard to distinguish precisely because their coordinates are equal.

Connect it to the computation

A JAX random function is deterministic given the same key and arguments. Calling it again with that key replays the result; it does not silently advance a hidden random state. Splitting produces a new key for a new draw.

The lines only connect coordinates to make equality visible. This is not a time series or a probability density, and three different values is distinct from statistical independence. The practical check is reproducibility: reuse for deliberate replay, split for the next stochastic operation.

Open full-size figureDownload data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Stage 2 of 5 · Build & Run the ReferenceStep-by-step code construction & verified execution output

Where to put & run this code: Work in phases/03-state/01-random-keys-without-surprises/code/main.py (or downloaded exercises/state-01.py / exercises/state-01.ipynb) and execute from the repository root with python3 phases/03-state/01-random-keys-without-surprises/code/main.py.

1 Prepare the inputs

Create main.py in your lesson workspace. Add this first block; use the environment from setup.

main.py · append build step 1
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

These explicit inputs define the case that the later checks will verify.

2 Build the computation

Append this block below the inputs in the same file.

main.py · append build step 2
# Step 2 — Build the computation: random.split returns two new keys.
# Create or split explicit PRNG key(s) (`key`) for reproducible randomness.
key = jax.random.key(42)
# Create or split explicit PRNG key(s) (`(key, sample_key)`) for reproducible randomness.
key, sample_key = jax.random.split(key)
# Sample deterministic random values into `a` using an explicit PRNG key.
a = jax.random.normal(sample_key, (3,))
# Sample deterministic random values into `repeated` using an explicit PRNG key.
repeated = jax.random.normal(sample_key, (3,))
# Create or split explicit PRNG key(s) (`(key, next_key)`) for reproducible randomness.
key, next_key = jax.random.split(key)
# Sample deterministic random values into `b` using an explicit PRNG key.
b = jax.random.normal(next_key, (3,))

Code building blocks in this step

  • jax.random.PRNGKey(seed) & jax.random.split(key) — Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.

random.split returns two new keys. noise_step returns the continuation key as part of its result so the caller can use fresh noise on the next transition.

3 Run and check the result

Append the checks, save main.py, and run python main.py from this folder using your course environment.

main.py · append build step 3
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print("Replay equal:", bool(jnp.array_equal(a, repeated)))
print("Fresh draw equal:", bool(jnp.array_equal(a, b)))
assert jnp.array_equal(a, repeated)
assert not jnp.array_equal(a, b)

Code building blocks in this step

  • jnp.array_equal(...) — Call jnp.array_equal with your updated parameters or inputs from this lesson's workspace.
  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

Compare the output to the expected result below before making the exercise change.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

# Step 2 — Build the computation: random.split returns two new keys.
# Create or split explicit PRNG key(s) (`key`) for reproducible randomness.
key = jax.random.key(42)
# Create or split explicit PRNG key(s) (`(key, sample_key)`) for reproducible randomness.
key, sample_key = jax.random.split(key)
# Sample deterministic random values into `a` using an explicit PRNG key.
a = jax.random.normal(sample_key, (3,))
# Sample deterministic random values into `repeated` using an explicit PRNG key.
repeated = jax.random.normal(sample_key, (3,))
# Create or split explicit PRNG key(s) (`(key, next_key)`) for reproducible randomness.
key, next_key = jax.random.split(key)
# Sample deterministic random values into `b` using an explicit PRNG key.
b = jax.random.normal(next_key, (3,))

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print("Replay equal:", bool(jnp.array_equal(a, repeated)))
print("Fresh draw equal:", bool(jnp.array_equal(a, b)))
assert jnp.array_equal(a, repeated)
assert not jnp.array_equal(a, b)

Run the example

python3 phases/03-state/01-random-keys-without-surprises/code/main.py
About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example
main.py · complete Python example
# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp
# Step 2 — Build the computation: random.split returns two new keys.
# Create or split explicit PRNG key(s) (`key`) for reproducible randomness.
key = jax.random.key(42)
# Create or split explicit PRNG key(s) (`(key, sample_key)`) for reproducible randomness.
key, sample_key = jax.random.split(key)
# Sample deterministic random values into `a` using an explicit PRNG key.
a = jax.random.normal(sample_key, (3,))
# Sample deterministic random values into `repeated` using an explicit PRNG key.
repeated = jax.random.normal(sample_key, (3,))
# Create or split explicit PRNG key(s) (`(key, next_key)`) for reproducible randomness.
key, next_key = jax.random.split(key)
# Sample deterministic random values into `b` using an explicit PRNG key.
b = jax.random.normal(next_key, (3,))
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print("Replay equal:", bool(jnp.array_equal(a, repeated)))
print("Fresh draw equal:", bool(jnp.array_equal(a, b)))
assert jnp.array_equal(a, repeated)
assert not jnp.array_equal(a, b)

Expected result: Replay equal: True; Fresh draw equal: False.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:45:26.921075+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Replay equal: True
Fresh draw equal: False
Replay equal: True
Fresh draw equal: False
PASS: state-01

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Stage 3 of 5 · Controlled ExperimentsPredict before running & test boundary conditions

Implement one random state transition

Implement one random state transition · Python experiment
# Experiment — Implement one random state transition: Replay compares complete state transitions.
def noise_step(key):
    # Create or split explicit PRNG key(s) (`(next_key, draw_key)`) for reproducible randomness.
    next_key,draw_key=jax.random.split(key)
    return next_key,jax.random.normal(draw_key,(2,))
# Create or split explicit PRNG key(s) (`root`) for reproducible randomness.
root=jax.random.key(9)
k1,n1=noise_step(root)
k2,n2=noise_step(k1)
# Create or split explicit PRNG key(s) (`(r1, replay1)`) for reproducible randomness.
r1,replay1=noise_step(jax.random.key(9))
r2,replay2=noise_step(r1)
assert jnp.array_equal(n1,replay1) and jnp.array_equal(n2,replay2)
assert jnp.array_equal(jax.random.key_data(k2),jax.random.key_data(r2))

Expected: Reconstructing the two steps reproduces both samples and the continuation key.

Replay compares complete state transitions. It catches lost continuation state, not only a repeated final number.

Keep identity under reordering

Keep identity under reordering · Python experiment
# Experiment — Keep identity under reordering: fold_in makes the assignment explicit.
# Create or split explicit PRNG key(s) (`base`) for reproducible randomness.
base=jax.random.key(11)
def by_id(ids):
    return jax.vmap(lambda i:jax.random.normal(jax.random.fold_in(base,i),()))(ids)
forward=by_id(jnp.array([3,8]))
reverse=by_id(jnp.array([8,3]))
assert jnp.array_equal(forward,reverse[::-1])

Expected: Noise follows the identifiers, not their batch positions.

fold_in makes the assignment explicit. Repeated IDs would deliberately reproduce a draw.

Stage 4 of 5 · Hands-On Exercise & CheckpointModify the code yourself, solve transfer tasks & test your understanding
Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/03-state/01-random-keys-without-surprises/code/main.py (or inside the Make it yours cell of exercises/state-01.ipynb) and run python3 phases/03-state/01-random-keys-without-surprises/code/main.py to verify your prediction before opening the reference solution.

Split a fresh seed into four keys, then use vmap to draw two normal values from each. Replay from the seed and compare.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jax.vmap(fn, in_axes=..., out_axes=...) — Vectorizes a single-example function across a batch axis without writing a Python loop.
  • jax.random.PRNGKey(seed) & jax.random.split(key) — Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.

2. Step-by-step implementation plan:

  1. Create or split explicit PRNG key(s) (keys) for reproducible randomness.
  2. Check tensor shape invariant: draw(7).shape == (4, 2)
Starter code scaffold · copy into your file & fill in the TODOs 
# Exercise solution: Split a fresh seed into four keys, then use vmap to draw two normal...
def draw(seed):
    # Create or split explicit PRNG key(s) (`keys`) for reproducible randomness.
    keys = jax.random.split(...)  # TODO: compute keys
    return ...  # TODO: return computed result
# Check tensor shape invariant: `draw(7).shape == (4, 2)`
assert draw(7).shape  # TODO: complete assertion check
assert jnp.array_equal(draw(7), draw(7))  # TODO: complete assertion check
Reveal the reference solution
Reference solution
# Exercise solution: Split a fresh seed into four keys, then use vmap to draw two normal...
def draw(seed):
    # Create or split explicit PRNG key(s) (`keys`) for reproducible randomness.
    keys = jax.random.split(jax.random.key(seed), 4)
    return jax.vmap(lambda k: jax.random.normal(k, (2,)))(keys)
# Check tensor shape invariant: `draw(7).shape == (4, 2)`
assert draw(7).shape == (4, 2)
assert jnp.array_equal(draw(7), draw(7))

Replay a noisy prediction

Practice

Draw one noise value per example for prediction 2x+12x+1. Reconstruct the same batch from the seed, then change only xx.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jax.random.PRNGKey(seed) & jax.random.split(key) — Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.

2. Step-by-step implementation plan:

  1. Create or split explicit PRNG key(s) (eps) for reproducible randomness.
Starter code scaffold · copy into your file & fill in the TODOs 
# Replay a noisy prediction (Practice): Holding noise fixed isolates the changed numerical input.
def noisy_prediction(seed,x):
    # Create or split explicit PRNG key(s) (`eps`) for reproducible randomness.
    eps = jax.random.normal(...)  # TODO: compute eps
    return ...  # TODO: return computed result
inputs = jnp.array(...)  # TODO: compute inputs
assert jnp.array_equal(noisy_prediction(4,inputs),noisy_prediction(4,inputs))  # TODO: complete assertion check
assert jnp.allclose(noisy_prediction(4,inputs+1)-noisy_prediction(4,inputs),2.,atol=1e-6)  # TODO: complete assertion check
Show a hint

Reuse the random construction for a controlled comparison, not accidentally across independent training steps.

Reveal reference solution and reasoning
Replay a noisy prediction · Python solution
# Replay a noisy prediction (Practice): Holding noise fixed isolates the changed numerical input.
def noisy_prediction(seed,x):
    # Create or split explicit PRNG key(s) (`eps`) for reproducible randomness.
    eps=jax.random.normal(jax.random.key(seed),x.shape)
    return 2*x+1+eps
inputs=jnp.array([0.,1.,2.])
assert jnp.array_equal(noisy_prediction(4,inputs),noisy_prediction(4,inputs))
assert jnp.allclose(noisy_prediction(4,inputs+1)-noisy_prediction(4,inputs),2.,atol=1e-6)

Holding noise fixed isolates the changed numerical input. That is a deliberately coupled experiment.

Diagnose key reuse in a loop

Challenge

Write the broken loop that draws with one key three times, demonstrate repeated rows, then repair it by threading the continuation.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jax.random.PRNGKey(seed) & jax.random.split(key) — Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.

2. Step-by-step implementation plan:

  1. Create or split explicit PRNG key(s) (root) for reproducible randomness.
  2. Sample deterministic random values into broken using an explicit PRNG key.
  3. Combine or mask array elements to form fixed.
  4. Check tensor shape invariant: fixed.shape==(3,2)
Starter code scaffold · copy into your file & fill in the TODOs 
# Diagnose key reuse in a loop (Challenge): The repeated-row symptom comes from key allocation.
# Create or split explicit PRNG key(s) (`root`) for reproducible randomness.
root = jax.random.key(...)  # TODO: compute root
# Sample deterministic random values into `broken` using an explicit PRNG key.
broken = jnp.stack(...)  # TODO: compute broken
assert jnp.array_equal(broken[0],broken[1])  # TODO: complete assertion check
current = ...  # TODO: compute current
rows = ...  # TODO: compute rows
for _ in range(3):
    current,noise = noise_step(...)  # TODO: compute current,noise
    rows.append(noise)
# Combine or mask array elements to form `fixed`.
fixed = jnp.stack(...)  # TODO: compute fixed
# Check tensor shape invariant: `fixed.shape==(3,2)`
assert fixed.shape  # TODO: complete assertion check
assert not jnp.array_equal(fixed[0],fixed[1])  # TODO: complete assertion check
Show a hint

The draw function does not mutate the key.

Reveal reference solution and reasoning
Diagnose key reuse in a loop · Python solution
# Diagnose key reuse in a loop (Challenge): The repeated-row symptom comes from key allocation.
# Create or split explicit PRNG key(s) (`root`) for reproducible randomness.
root=jax.random.key(5)
# Sample deterministic random values into `broken` using an explicit PRNG key.
broken=jnp.stack([jax.random.normal(root,(2,)) for _ in range(3)])
assert jnp.array_equal(broken[0],broken[1])
current=root
rows=[]
for _ in range(3):
    current,noise=noise_step(current)
    rows.append(noise)
# Combine or mask array elements to form `fixed`.
fixed=jnp.stack(rows)
# Check tensor shape invariant: `fixed.shape==(3,2)`
assert fixed.shape==(3,2)
assert not jnp.array_equal(fixed[0],fixed[1])

The repeated-row symptom comes from key allocation. The last inequality checks this fixed demonstration, not every possible random stream.

Check your understanding

What happens if the same key is used twice for the same random operation?

Choose one answer

Diagnose the result

The repeated-row symptom comes from key allocation. The last inequality checks this fixed demonstration, not every possible random stream.

Stage 5 of 5 · Keep Your EvidenceRecord your verification receipt & carry takeaways forward

Carry forward

  • Replay compares complete state transitions. It catches lost continuation state, not only a repeated final number.
  • fold_in makes the assignment explicit. Repeated IDs would deliberately reproduce a draw.

Keep your evidence

Download evidence template

Keep a key-ownership tree, two-step replay including the final key, the per-ID reorder check, and the broken/repaired repeated-noise loop.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references