Skip to content

← JAX Foundations step 13 of 25

Easy Primitives

PRNGKey and Split

Jump to the task

A key makes a random draw repeatable

JAX's functional random API takes a key explicitly. A key is an input to the calculation: sampling does not change it. Calling the same sampler with the same key, shape and other settings repeats the result. You do not advance to a new draw merely by calling the function again.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
first = jax.random.normal(key, (3,))
repeated = jax.random.normal(key, (3,))
# first and repeated contain the same values.

To request separate draws, start from a fresh base key and split it:

base = jax.random.PRNGKey(1)
key1, key2 = jax.random.split(base, 2)
first = jax.random.normal(key1, (3,))
second = jax.random.normal(key2, (3,))
result = jnp.stack([first, second])  # shape (2, 3)

Split deterministically derives child keys for separate pseudorandom streams. Use each child once, and stop using the parent after splitting. Independent draws can still happen to contain equal values; independence does not mean every value must differ. The repeated-key example above demonstrates deliberate replay, not a way to get independent draws.

Compare key strategies

Switch between reuse, splitting and folding in an index below. The diagram shows repeated scalar sampling calls, with illustrative values rather than actual JAX output. Splitting each step carries a new key forward; fold_in derives a key from the base and a step index. Neither pattern is needed just to draw several values in one array.

Loading visualization…

Keys and sample shapes are different

This exercise uses the legacy PRNGKey API with the default Threefry generator. Its key has shape (2,) and dtype uint32; splitting into two gives keys with shape (2, 2). Those numbers describe keys, not samples.

Modern JAX also provides typed keys through jax.random.key(seed). A typed key is scalar, so splitting it into two gives shape (2,). Both representations work with split and normal. Do not assume all JAX keys are uint32 arrays with a final dimension of length two.

normal(key, (n,)) requests n values. Omitting the shape requests one scalar. You do not need n different keys to draw a vector of n values; here you need two keys because the task asks for two separate vectors.

Problem

Implement two_random_draws(seed, n):

  1. Create the base key with jax.random.PRNGKey(int(seed)).
  2. Split it into exactly two keys, keeping their returned order.
  3. Draw one float32 standard-normal vector of shape (int(n),) per key.
  4. Stack them in that order to return shape (2, int(n)).

Seed is a whole Python number in [0, 2**32), and n is a nonnegative whole Python number. Both arrive as floats; convert them to int. For n = 0, return shape (2, 0). Use the runtime's default Threefry sampling configuration so the results match the expected stream.

A single draw of shape (2, n) from the unsplit base key is a different stream, even though its shape and distribution look right. The tests check the reproducible values, not just the distribution.

To use this particular interface with JIT, mark both arguments static: jax.jit(two_random_draws, static_argnames=("seed", "n")). int() cannot convert a traced value into a static one. A redesigned interface could accept a dynamic key, but the sample shape still needs to be known during tracing.