A pseudo-random number generator (PRNG) deterministically produces values that appear random from a seed. Stateless systems require the seed or key to be passed explicitly instead of hiding it in global state. section 3.5.3
The key implementation ideas are:
- R and NumPy-style calls can depend on hidden global generator state; setting a seed makes repeated operations reproducible. section 3.5.3
- JAX uses explicit PRNG key objects, passed as the first argument to random operations. section 3.5.3
- Reuse a key to reproduce the same result; split it to obtain a new key and different random values. section 3.5.3
Reproducible noise and JAX values
An R function can set a supplied seed before adding uniform noise, so identical inputs and seeds produce identical outputs. In JAX, create a key with jax$random$key(123L); repeated normal calls with that key match, while jax$random$split() supplies a new key for different values. section 3.5.3
seed_key <- jax$random$key(123L) jax$random$normal(seed_key, shape = shape(3)) Array([1.6359469 , 0.8408094 , 0.02212393], dtype = float32) new_seed_key <- jax$random$split(seed_key, num = 1L)@py[0] #1 jax$random$normal(new_seed_key, shape = shape(3)) Array([-0.49093357, -0.9478693 , -1.775197 ], dtype = float32)
Inspect the complete JAX key-reuse and key-splitting example. section 3.5.3
