Skip to content

← JAX Stochasticity step 18 of 25

Medium Primitives

Batched Sampling via vmap

Jump to the task

Why this matters

JAX random functions take an explicit key. Reusing a key with the same sampler, shape and configuration repeats the same values. To give each row its own random draw, split a base key into row keys and map the sampler over them.

This pattern is useful when the per-row calculation will later grow beyond sampling. vmap adds the batch axis without requiring you to rewrite that calculation to handle a whole batch.

Worked mini-example

import jax

base_key = jax.random.PRNGKey(0)
keys = jax.random.split(base_key, 2)
samples = jax.vmap(lambda key: jax.random.normal(key, (3,)))(keys)
# keys.shape == (2, 2): two legacy keys, each holding two uint32 words.
# samples.shape == (2, 3): one length-three sample per key.

vmap removes the leading axis from the sampler's view: each invocation receives one key, then returns a vector of three values. Stacking those vectors creates the two rows. The keys produce pseudorandom streams designed for independent sampling; individual values can still coincide.

Problem

Implement vmap_normal_batch(seed, batch_size, n):

  1. Create jax.random.PRNGKey(int(seed)).
  2. Split it with jax.random.split(base_key, int(batch_size)).
  3. Apply jax.random.normal(key, (int(n),)) directly to each split key using jax.vmap.

seed is a nonnegative integer-valued scalar. batch_size and n are nonnegative integer-valued counts. They may arrive as floats. With the default float32, Threefry random configuration, return a float32 array of shape (int(batch_size), int(n)). Zero counts retain both axes: for example, batch_size=0 and n=3 gives shape (0, 3).

The fixtures check the values from this exact split-and-draw procedure. Drawing one array with jax.random.normal(base_key, (batch_size, n)) is also a valid way to generate a batch, but generally produces different values and does not satisfy this exercise's seeded-output contract.

Common pitfalls

  • Passing the same key to every mapped draw repeats the same row. Key splitting, rather than vmap itself, supplies different streams.
  • An extra split inside the sampler changes the keys and expected values. Draw directly from the keys produced by the initial split.
  • All three Python conversions need concrete values under JIT. Mark seed, batch_size and n static when wrapping this function in jax.jit; int(...) cannot make a dynamic tracer concrete.
  • The key array's final dimension of two is storage for a legacy key, not a sample dimension. Typed keys created with jax.random.key have different array shapes. This exercise and diagram use PRNGKey.
  • A Python loop over the same keys can produce the same values; vmap expresses the batch transformation directly. Neither syntax alone guarantees a speedup for every workload.

The diagram follows the worked example: each mapped invocation receives a legacy key of shape [2] and returns a sample of shape [3].

Loading visualization…