Determinism means getting the same result when a computation receives the same inputs and explicit random state. JAX makes random state explicit: calls using the same PRNG key return the same values. section 3.5.3
To manage random values in JAX:
- Create a PRNG key from an integer seed and pass it to random operations. section 3.5.3
- Split a key to create a different key when you need new random values. section 3.5.3
- In programs that combine R, Python, NumPy, and Keras, use keras3::set_random_seed() to set seeds across those systems; set.seed() affects only R. section 3.5.3
seed_key <- jax$random$key(123L) jax$random$normal(seed_key, shape = shape(3)) jax$random$normal(seed_key, shape = shape(3)) Array([1.6359469 , 0.8408094 , 0.02212393], dtype = float32) Array([1.6359469 , 0.8408094 , 0.02212393], dtype = float32)
The same JAX key passed to two normal-random calls produces identical arrays. 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)
Splitting a key provides a different key and produces different random values. section 3.5.3
