A pseudorandom number generator produces random-like values from a seed or internal state. In JAX, random operations use explicit keys rather than hidden global state, so repeating a call with the same key returns the same value. section 3.5.3
To manage random values in JAX:
- Create a key from an integer seed and pass it as the first argument to random operations. section 3.5.3
- Split an existing key to get a different key when you need new random values; the same sequence of splits is deterministic. section 3.5.3
- Manage keys outside functions that need fresh values. This takes more work than relying on NumPy’s global state, but supports deterministic, vectorizable JAX code. section 3.5.3
Listing 3.30 Using a random seed in Jax
>>> seed_key = jax.random.key(123) >>> jax.random.normal(seed_key, shape=(3,)) Array([-0.1470326, 0.5524756, 1.648498 ], dtype=float32) >>> jax.random.normal(seed_key, shape=(3,)) Array([-0.1470326, 0.5524756, 1.648498 ], dtype=float32)
The repeated calls use the same key and produce identical random tensors. section 3.5.3
