JAX does not use a global random-number state. Instead, each random operation receives a key: a special array structure created from an integer seed. Pass the key as the first argument to a call such as `jax.random.normal(key, shape=(3,))`. 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 two calls use the same key and return the same values. A key does not silently advance when used. section 3.5.3
A practical rule of thumb:
- Use the same key when you need the same random values again. section 3.5.3
- Use `jax.random.split()` to derive a new key when you need different values. The same sequence of splits produces the same final key. section 3.5.3
>>> 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] #1 >>> jax.random.normal(new_seed_key, shape=(3,)) Array([ 0.5362355, -1.1920372, 2.450225 ], dtype=float32)
Splitting the original key changes the values returned by `random.normal()` without changing global state. section 3.5.3
Managing keys is more work than relying on NumPy’s global PRNG state, but it keeps random operations stateless. That supports repeatable results while allowing the JAX compiler to vectorize code. section 3.5.3
