--- jupytext: formats: md:myst text_representation: extension: .md format_name: myst format_version: 0.13 jupytext_version: 1.16.4 kernelspec: display_name: Python 3 language: python name: python3 --- (jax-101-random)= # Pseudorandom numbers In most numerical libraries, drawing a random number means consulting a hidden generator that updates itself between calls. JAX has no such generator, and no hidden update. Randomness in JAX is driven by **keys**: ordinary immutable array values that you pass around explicitly. Every function in {mod}`jax.random` is a pure, deterministic function of its key argument: call it twice with the same key and you get the same number twice, by design. ## Why not a global generator? NumPy's `numpy.random` is the familiar stateful design: a global generator, seeded once, silently advanced by every sampling call: ```{code-cell} import numpy as np np.random.seed(0) print(np.random.uniform()) print(np.random.uniform()) # different: the hidden state advanced ``` The problem is that "the hidden state advanced" makes your results depend on the exact order and number of sampling calls everywhere in your program: ```{code-cell} np.random.seed(0) def bar(): return np.random.uniform() def baz(): return np.random.uniform() def foo(): return bar() + 2 * baz() print(foo()) ``` This value is only reproducible because NumPy promises to run `bar()` before `baz()`. Such sequencing promises are what JAX needs to avoid: an optimizing compiler should be free to reorder and parallelize work across devices. JAX needs random number generation that is **reproducible**, **parallelizable**, and **vectorizable**, which rules out sampling functions that secretly read and write shared state. JAX's solution isn't simply to make the generator state an explicit argument that you shuttle in and out of functions. Instead, sampling is a pure function of a key, and no updated state comes back out. ## Keys are values Create a key from an integer seed with {func}`jax.random.key`: ```{code-cell} import jax from jax import random key = random.key(42) key ``` A key is a rank-0 array (shape `()`) with a special element type: its JAX type is `key[]`, where `key` names the default Threefry PRNG implementation. ```{code-cell} print(jax.typeof(key)) ``` Passing a key to a sampling function doesn't modify it or "use it up" in any physical sense; the sample is a deterministic function of the key: ```{code-cell} print(random.normal(key)) print(random.normal(key)) # same key, same value — necessarily ``` This means reproducibility is automatic: results depend only on the key values your program constructs, never on execution order, call counts, or which device ran what. The flip side is that *distinct* random numbers require *distinct* keys, which leads to the one rule of JAX randomness: **Never reuse a key (unless you want identical outputs).** Feeding the same key to two different samplers produces correlated results, depriving your program of life-giving chaos. ## Deriving new keys To get fresh keys, derive them from a key you already have. {func}`jax.random.split` deterministically derives new keys, each of which can be used to generate statistically independent samples: ```{code-cell} key = random.key(42) key, subkey = random.split(key) print(random.normal(subkey)) ``` `split` produces as many keys as you ask for in one shot: ```{code-cell} key = random.key(42) subkeys = random.split(key, num=4) [float(random.normal(k)) for k in subkeys] ``` And {func}`jax.random.fold_in` derives a new key from a key and an integer, which is ideal for generating a per-step or per-example key without threading a key through your loop: ```{code-cell} key = random.key(42) for step in range(3): step_key = random.fold_in(key, step) print(f"step {step}: {random.normal(step_key)}") ``` Note the shape of this pattern: every `step_key` is derived directly from one parent key, not from its predecessor. The next section explains why. ### Keep the key tree wide, not deep Your program's keys form a tree, rooted at the seed, growing by either `split` or `fold_in` operations. A bad pattern is to grow that tree as a long chain, splitting off each step's key from the previous step's: ```python for step in range(num_steps): key, subkey = random.split(key) # each key derived from the last: avoid! ... ``` Prefer wide trees instead: all step keys hanging off one parent, via a single `split(key, num_steps)` or via `fold_in(key, step)` as above. The chained version has two problems, one computational and one statistical: - **It serializes.** Each key depends on the previous one, so a chain of a million steps means a million sequential hash applications. Wide key derivation is a single batched operation, free to vectorize and parallelize. - **It courts collisions.** For a *fixed* key, the PRNG's underlying hash is a pseudorandom *permutation* of its input, so the keys produced by one `split` (or by `fold_in` over distinct integers) are guaranteed distinct. But viewed as a function *of the key*, the hash is not a permutation: it behaves like a random function. Every derivation hop is therefore an independent chance for two keys in your tree to coincide, and over a long chain in the default 64-bit key space, collision probability accumulates toward the birthday bound. A collision means identical random streams from the point of collision onward. A handful of chained splits is harmless; the collision math only bites at scale, and plenty of correct code splits a key a few times in sequence. But for anything proportional to the length of training or the size of a dataset, derive keys widely from a common parent. ## No sequential equivalence NumPy guarantees that sampling N numbers one at a time yields the same sequence as sampling N at once. JAX deliberately makes no such promise: ```{code-cell} key = random.key(42) subkeys = random.split(key, 3) print("individually:", np.stack([random.normal(k) for k in subkeys])) key = random.key(42) print("all at once: ", random.normal(key, shape=(3,))) ``` Sequential equivalence would impose the kind of ordering constraint JAX's design exists to avoid. Without it, samples drawn from different keys have no ordering relationship at all, so generation can be freely vectorized and sharded. Since keys are just arrays, they compose with everything else in JAX. You can `vmap` a sampler over a batch of keys: ```{code-cell} jax.vmap(random.normal)(subkeys) ``` With the default PRNG implementation, this is *exactly* equivalent to calling `random.normal` on each key separately; vectorizing over keys doesn't change the values. ```{note} These docs use the typed keys created by {func}`jax.random.key`. You may also encounter older code using {func}`jax.random.PRNGKey`, which produces a raw `uint32` array — it still works, but it's easy to misuse (nothing stops you from doing arithmetic on it) and it doesn't record which PRNG implementation it belongs to. Prefer `jax.random.key` in new code, and convert at boundaries with {func}`jax.random.key_data` and {func}`jax.random.wrap_key_data` when interfacing with systems that need raw arrays. See the [typed PRNG keys JEP](https://docs.jax.dev/en/latest/jep/9263-typed-keys.html) for the full story. ``` ## Design and implementations JAX's PRNG is a counter-based Threefry hash combined with a functional splitting model, chosen so that generation has no sequencing constraints at all. See the [PRNG design JEP](https://docs.jax.dev/en/latest/jep/263-prng.html) for the design rationale. Threefry is the default of several available implementations. Alternatives (selected per-key via the `impl` argument to {func}`jax.random.key`, or globally via the `jax_default_prng_impl` config) trade off generation speed on TPUs, shardability and sharding-invariance, bit-for-bit identical results across platforms, and exact `vmap`-over-keys semantics. The {mod}`jax.random` module documentation has the full comparison table; the default is the right choice unless PRNG generation shows up in your profiles. {mod}`jax.random` itself offers samplers for a wide range of distributions — uniform, normal, categorical, permutations, and many more — all taking a key as their first argument. ## Next steps Keys handle randomness while keeping every function pure. The remaining topic is state: values that evolve as a program runs, and genuine in-place mutation. It's covered in {ref}`jax-101-state`.