TopolearnJAX's RNG is reproducible and parallel-safe only if you handle keys correctly; reusing a key is a silent bug.
NumPy's np.random keeps hidden global state, which breaks the purity that jit and vmap rely on. JAX instead makes randomness explicit: every random function takes a key and is a pure function of it. The same key always produces the same numbers.
key = jax.random.key(42)
x = jax.random.normal(key, (3,))The rule: never reuse a key for two different random draws you expect to be independent. Instead split it:
key, sub = jax.random.split(key)gives a fresh subkey to consume and a new key to carry forward.jax.random.split(key, n)returnsnkeys at once, handy withvmap(e.g. one dropout key per example).jax.random.fold_in(key, step)derives a key from an integer such as the step number or device index.
Inside a training loop you thread the key through the step function just like parameters: take it in, split, use one part, and return the other.
jax.random.key creates a typed key (dtype key<fry>), the recommended API. Older code uses jax.random.PRNGKey(seed), which returns a raw uint32 array of shape (2,); it still works, but typed keys catch misuse such as doing arithmetic on a key.
Random functions need static shapes: jax.random.uniform(key, (n,)) requires n to be a concrete Python int under jit.
Example
import jax
key = jax.random.key(0)
a = jax.random.normal(key, (3,))
b = jax.random.normal(key, (3,))
print((a == b).all()) # True: same key, same numbers (a bug if unintended)
key, k1, k2 = jax.random.split(key, 3)
print((jax.random.normal(k1, (3,)) == jax.random.normal(k2, (3,))).any()) # False
def dropout(key, x, rate=0.5):
keep = jax.random.bernoulli(key, 1 - rate, x.shape)
return jax.numpy.where(keep, x / (1 - rate), 0.0)
keys = jax.random.split(key, 4) # one key per example
xs = jax.numpy.ones((4, 5))
print(jax.vmap(dropout)(keys, xs).shape) # (4, 5)
step_key = jax.random.fold_in(key, 100) # derive from a step number
print(step_key.dtype) # key<fry>Remember
- Random functions are pure functions of an explicit key
- Split keys before every independent use; never reuse one
- Prefer typed keys from
jax.random.keyover legacyPRNGKey
Watch out for
- Using the same key for weight init in every layer, which makes the layers' weights identical or correlated.