Skip to content

← JAX Foundations step 15 of 25

Medium Primitives

Deterministic Batch via vmap+split

Jump to the task

Give each row its own key

Suppose you want four rows of random noise, with three values in each row. First split a base key into four child keys. Then define a function that draws three standard-normal values from one key. vmap applies that function to the keys and stacks the resulting rows.

This keeps each row's random input explicit. The draw function describes one row; it does not need to handle the batch dimension itself. It is useful when a larger per-example calculation also needs its own key.

Worked mini-example

import jax, jax.numpy as jnp

base = jax.random.PRNGKey(0)
batch_size = 4

keys = jax.random.split(base, batch_size)   # shape (4, 2)
def draw(k):
    return jax.random.normal(k, (3,))       # one length-3 sample per key

batch = jax.vmap(draw)(keys)
# batch.shape == (4, 3): each row is an independent N(0,1) sample

Common pitfalls

  • Whole-array sampling is valid too: jax.random.normal(base, (batch_size, n)) also generates independent standard-normal entries. It generally produces different values from split-then-vmap: JAX does not promise sequential equivalence between these sampling strategies. This exercise specifically checks the split-key construction.
  • Forgetting to vmap: a Python list comprehension over keys would work but unrolls inside jit.
  • Wrong n cast: jax.random.normal(k, (n,)) requires n as a Python int. Cast with int(n) outside the vmap'd draw if needed (the closure captures it). If you later wrap the outer function in jit, these shape arguments must be static; int() cannot make a dynamic tracer concrete.
  • Extra splitting changes the random stream: splitting inside a mapped function can be valid, but is not the construction specified here.

The diagram follows the worked example: four legacy PRNGKey values of shape [2] produce four length-three samples. Typed keys have a different storage shape; this exercise uses the legacy API shown above.

See JAX's sampling and sequential-equivalence guide.

Loading visualization…

Problem

Implement deterministic_batch(seed, batch_size, n):

  • Create a base key from seed using jax.random.PRNGKey(int(seed)).
  • Split it into batch_size sub-keys with jax.random.split.
  • Define a per-sample draw function that draws n values from N(0, 1).
  • Apply the draw function over all sub-keys using jax.vmap.
  • Return a 2-D array of shape (batch_size, n).

All three arguments are whole Python numbers supplied as floats. Seed is in [0, 2**32); batch_size and n are nonnegative. Convert them to int. Use float32 draws and the default Threefry configuration used by the exercise runtime. Preserve the order returned by split.

A zero batch size returns shape (0, n). Zero values per row returns shape (batch_size, 0). Repeating the same arguments and configuration reproduces the same batch; a repeat call does not advance hidden state.

To JIT this float-based interface, mark all three arguments static: jax.jit(deterministic_batch, static_argnames=("seed", "batch_size", "n")). Seed is static here because int(seed) is a Python conversion. An interface that accepts a key could keep that key dynamic, while the requested output shape still needs to be known during tracing.

Output tests verify the chosen random stream and shape. They cannot distinguish vmap from another implementation that computes the same rows.

Example output shapes:

  • deterministic_batch(1, 3, 4) returns a (3, 4) array: three rows of four independent standard-normal values, one row per batch element.
  • deterministic_batch(99, 5, 1) returns a (5, 1) array: five scalars, each drawn from N(0, 1) via its own split-derived key.