A JAX key is an explicit value representing pseudorandom-number-generator state. Random functions receive the key as an argument instead of reading or mutating hidden global state, keeping randomness compatible with pure functions, transformations, reproducibility, and parallel execution. section 9.0 section 9.2.3
Use keys according to three practical rules:
- Create one with random.PRNGKey(seed), using a 32- or 64-bit integer seed. section 9.1.2
- Pass a key explicitly to each random operation; unlike NumPy's hidden state, JAX random functions do not update the key. section 9.2.3
- Split a key before separate random operations, and discard the old key after using it for generation, split(), or fold_in(). Reusing one key reproduces the same values and can entangle operations. section 9.1.3 section 9.4
- Request multiple values with one key and a shape when a single random tensor is enough; split into subkeys when later operations need separate random streams. section 9.2.3
Generate independent random values
from jax import random key = random.PRNGKey(42) key1, key2 = random.split(key, num=2) vals = random.normal(key1, shape=(3, 5)) more_vals = random.normal(key2, shape=(3, 5)) The two subkeys deterministically produce different random arrays. Calling random.normal(key, ...) twice instead would produce the same array twice. For repeated indexed work, random.fold_in(key, i) creates a new key from the original key and an integer. section 9.2.3 section 9.4
Drive a random augmentation
An augmentation pipeline can split one incoming key into one key for choosing an augmentation and another for the selected operation. A noise augmentation consumes its key through random.normal(), while a horizontal flip can accept the same function signature even when it does not use its key. This keeps branching randomness explicit and avoids reusing a key. section 9.1.3
Listing 9.16 Generating different random values in JAX
from jax import random key = random.PRNGKey(42) #1 key >>> Array([ 0, 42], dtype=uint32) key1, key2 = random.split(key, num=2) #2 key1 #3 >>> Array([2465931498, 3679230171], dtype=uint32) key2 #3 >>> Array([255383827, 267815257], dtype=uint32) vals = random.normal(key1, shape=(3,5)) #4 more_vals = random.normal(key2, shape=(3,5)) #5 vals #6 >>> Array([[-0.95032114, 0.89362663, -1.9382219 , -0.9676806 , -0.3920417 ], >>> [ 0.6062024 , 0.37990445, 0.30284515, 1.3282853 , 1.1882905 ], >>> [-0.95166105, 1.2611355 , 0.41334143, -0.7831721 , 0.09786294]], >>> dtype=float32) more_vals #6 >>> Array([[ 0.7679137 , 0.46966743, -1.4884446 , -1.155719 , -1.2574353 ], >>> [-0.66447204, -1.3314192 , 1.5208852 , -0.55124223, -0.3213504 ], >>> [ 1.1168208 , -0.4216216 , 0.32398054, 1.3500887 , -0.22909231]], >>> dtype=float32) #1 Generates a key with some seed #2 Splits a key into two keys #3 The two keys are different. #4 Generates random normal values with the first key #5 Generates some more random normal values with the second key #6 Produces different numbers
Inspect the complete split-and-generate listing, including the distinct arrays produced from two subkeys. section 9.2.3
