PRNG Keys and Splitting

A PRNG key (Pseudo-Random Number Generator key) is a small array that encodes the random state. You create your first key from an integer seed, then produce new keys by splitting.

import jax
import jax.numpy as jnp

# Create a starting key from a seed
key = jax.random.key(0)
print(key)  # a special key array

# WRONG: reusing the same key gives the same numbers
print(jax.random.normal(key))  # 0.18784384
print(jax.random.normal(key))  # 0.18784384 — identical!

# RIGHT: split the key to get new, independent keys
key, subkey = jax.random.split(key)
print(jax.random.normal(subkey))  # -1.2515389

key, subkey = jax.random.split(key)
print(jax.random.normal(subkey))  # -0.5841975 — different!

jax.random.split(key) consumes one key and produces two new, independent keys. The standard pattern is to keep one for future splitting and use the other ("subkey") immediately. You can also split into more than two keys:

# Split into many keys at once
key = jax.random.key(42)
keys = jax.random.split(key, num=5)
print(keys.shape)  # (5,) — five independent keys

# Use each key for a different purpose
samples = jax.vmap(jax.random.normal)(keys)
print(samples)  # 5 independent random numbers

The key hierarchy forms a tree: you start with one key, split it into children, split those into grandchildren, and so on. Every path through the tree produces statistically independent random streams. This is mathematically guaranteed by the Threefry hash function that JAX uses under the hood.

A note on jax.random.key vs jax.random.PRNGKey: The modern API is jax.random.key(seed), which returns a typed key array. The older jax.random.PRNGKey(seed) still works but returns a raw uint32 array. Use key() for new code.

# Modern API (recommended)
key = jax.random.key(0)

# Legacy API (still works)
key_legacy = jax.random.PRNGKey(0)

# Both produce valid keys for all jax.random functions