Skip to content

← JAX Stochasticity step 19 of 25

Hard Primitives

Random Walk via lax.scan

Jump to the task

Build a path one random step at a time

A random walk starts at a position and adds a random change at each step. Here the initial position is zero. Each change is a standard normal draw multiplied by step_std, so step_std controls the size of the changes, not the standard deviation of the entire path.

Carry two things through lax.scan: the current position and a PRNG key for future draws. Split the key, use one child for this step's draw, and keep the other for the next step. Save the updated position as the scan output. This separates the state needed to continue from the history you want to return.

Worked mini-example

This two-step walk starts at zero with unit step standard deviation. The first split key is carried forward; the second supplies the next draw. This order is part of the exercise's reproducible stream.

import jax.random as jr
from jax import lax

def step(carry, _):
    key, pos = carry
    keys = jr.split(key)
    key, draw = keys
    dx = jr.normal(draw)
    pos = pos + dx
    carry = (key, pos)
    return carry, pos

key = jr.PRNGKey(0)
init = (key, 0.)
final, path = lax.scan(
    step, init, None,
    length=2
)
# path has shape (2,)

With the default Threefry PRNG and float32 draws, the positions are approximately [-2.4425, -3.6999]. Each step carries the fresh key and updated position forward, and saves only the position in path. The initial position zero is excluded: two steps produce two positions, not three. These are positions, not the raw changes: the second position adds the second change to the first position.

The same seed, split order and sampling configuration reproduce the same path. Drawing a whole vector from the original key or swapping the two child keys can define a valid random walk, but gives a different path and does not meet this exercise's reproducibility contract.

Common pitfalls

  • Reusing a key. Drawing with the same key and distribution parameters repeats the same delta. Split at each step and carry the first child key forward; use the second child for the draw.
  • Counting steps. You can pass xs=None, length=n, as above, or scan over a dummy array such as jnp.zeros(n). The latter is optional, not a requirement of scan.
  • Changing the carry structure. Keep (key, position) throughout this example. A consistently structured list can also be a valid pytree; switching from a tuple to a list during the scan is an error.
  • Dynamic length under JIT. The output length must be known when tracing. If wrapping this exercise in jit, its Python int conversions and output shape need static seed and n_steps arguments, or a redesigned interface accepting a key and static length.

For the supplied interface, JIT can be used like this after defining random_walk:

import jax

compiled = jax.jit(random_walk, static_argnames=("seed", "n_steps"))
path = compiled(0.0, 2.0, 1.0)

Marking seed static is needed here because the interface calls int(seed), not because JAX requires every random key or seed to be static. Changing these static arguments can trigger compilation again. A key-based interface avoids making the seed static, but the output length still needs to be known during tracing.

Problem

Implement random_walk(seed, n_steps, step_std):

The seed and n_steps are nonnegative whole Python numbers supplied as floats. Use int(seed) to create the PRNGKey and int(n_steps) for the scan length. step_std is a finite, nonnegative scalar.

At each step, split as key, subkey = jax.random.split(key). Carry key forward, draw a scalar standard normal using subkey, multiply by step_std, and add the change to the current position. Use the default Threefry configuration and float32 sampling used by the exercise runtime.

Return a float32 array of shape (n_steps,) with the position after each step. Exclude the initial zero. Zero steps returns an empty array; zero step_std returns an all-zero path. Output tests check the reproducible values, but cannot enforce the use of lax.scan.