A stateless system makes each function depend only on explicit inputs: it neither reads hidden global state nor quietly preserves values between calls. This makes repeated calls predictable and allows independent execution, which is valuable for parallel and distributed computation. section 3.5 section 3.5.3
Explicit state in a JAX training step
A training function receives model parameters explicitly and returns their replacements: `loss, W, b = training_step(inputs, targets, W, b)`. The function computes gradients from `(W, b)`, updates both values, and returns the loss plus updated parameters. Nothing about the parameters is stored implicitly between calls. section 3.5.4
Use these rules when converting stateful code:
- Pass mutable or evolving values—such as parameters, optimizer variables, metric variables, or random keys—as function arguments. section 3.5.4 section 7.4.5
- Return every changed value so the caller owns the next state. section 3.5.4
- Pass a JAX PRNG key to every random operation; reuse a key for reproducible output and split it deterministically for new values. 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 random call visibly produces the same array when given the same key, making determinism and explicit randomness concrete. section 3.5.3
