Skip to content

← JAX Foundations step 14 of 25

Medium Primitives

PRNGKey fold_in

Jump to the task

Derive a key from a step number

Imagine returning to training step 7 to reproduce its random noise. jax.random.fold_in(base_key, 7) derives a key from the base and that identifier. You can derive it directly without generating keys for steps 0 through 6 first. Keep the base fixed and change the identifier to obtain a stream for each step.

import jax
import jax.numpy as jnp

base = jax.random.PRNGKey(0)
step7 = jax.random.fold_in(base, 7)
replay7 = jax.random.fold_in(base, 7)
jnp.array_equal(step7, replay7)  # True

noise = jax.random.normal(step7, (3,))
replay = jax.random.normal(replay7, (3,))
# noise and replay contain the same values.

This deliberately replays a draw. For fresh noise, use a different identifier or split the derived key for separate uses. Repeating the same base and identifier does not create a new stream. Reproducibility also depends on the sampler, shape and sampling configuration.

fold_in and split solve different bookkeeping problems

Fold in an identifier when you want to address a stream directly, such as a training step or example number. Split when you want child keys to pass to separate calculations. You can split repeatedly as a loop runs; you do not need to know its final length in advance.

For this task, derive every row's key from the original base. Do not replace base with the previous row's key: that would create a chain of derivations and change the expected samples. Computing rows in a different order does not change a row's key when its identifier stays the same. Extending the number of rows preserves the existing rows when the sample shape and configuration are unchanged.

Integer identifiers can be dynamic

Fold_in accepts a scalar 32-bit integer identifier, including a traced JAX integer. Do not call Python int() on a traced loop index. A Python range already supplies integer indices; a JAX loop can pass its integer index directly. A vector of identifiers needs mapping, such as vmap.

A static Python loop works under JIT, but is unrolled during tracing. For long sequences, lax.map, scan or vmap can avoid a large unrolled program. Keep the same base-and-index rule when changing the loop style.

Problem

Implement fold_in_per_step(seed, n_steps, n):

  1. Create base_key with jax.random.PRNGKey(int(seed)).
  2. For each k from 0 through int(n_steps) - 1, derive jax.random.fold_in(base_key, k).
  3. Draw a float32 standard-normal vector of length int(n) from that key.
  4. Return the rows in increasing k order, with shape (int(n_steps), int(n)).

All arguments are whole Python numbers supplied as floats. Seed is in [0, 232), n_steps is in [0, 232), and n is nonnegative. Convert these interface values to int. Use the runtime's default Threefry configuration. Zero steps returns shape (0, n), and zero samples per step returns shape (n_steps, 0).

To JIT this interface, mark all three arguments static: jax.jit(fold_in_per_step, static_argnames=("seed", "n_steps", "n")). This requirement comes from the Python conversions and output shape, not from a requirement that fold_in's identifier be static.