jax.random.split deterministically derives a new JAX random key from an existing key. Use it when a random operation needs a fresh key without relying on global random state. section 3.5.3
Derive and use a fresh key
>>> seed_key = jax.random.key(123) >>> jax.random.normal(seed_key, shape=(3,)) Array([-0.1470326, 0.5524756, 1.648498 ], dtype=float32) >>> new_seed_key = jax.random.split(seed_key, num=1)[0] >>> jax.random.normal(new_seed_key, shape=(3,)) Array([ 0.5362355, -1.1920372, 2.450225 ], dtype=float32) Keep the original key and explicitly derive another key for the next random operation. section 3.5.3
Practical rules:
- Pass keys explicitly as the first argument to JAX random functions. section 3.5.3
- Expect repeated use of the same key to reproduce the same random value; split when a fresh value is required. section 3.5.3
