A stateless function receives changing computational state through its arguments and returns updated state explicitly. Nothing essential is hidden in mutable external variables. section 3.5.4
A stateless linear-model update
A JAX model accepts inputs, weights, and bias as arguments. A training step computes loss and gradients, updates the weights and bias, and returns the loss together with their new values. Repeated calls therefore carry the complete evolving state explicitly. section 3.5.4
A stateless training function typically follows this sequence:
- Accept model variables and any other changing values as arguments. section 3.5.4
- Compute predictions, loss, and gradients from the supplied state. section 3.5.4
- Return updated variables rather than modifying hidden state; include optimizer, trainable, and non-trainable variables when they participate in training. section 7.4.2
learning_rate <- 0.1
training_step <- jax$jit(\(inputs, targets, W, b) { #1
.[loss, grads] <- grad_fn(list(W, b), inputs, targets) #2
.[grad_wrt_W, grad_wrt_b] <- grads
W <- W - grad_wrt_W * learning_rate #3
b <- b - grad_wrt_b * learning_rate #3
tuple(loss, W, b) #4
})
This listing shows the complete shape of a stateless JAX update: receive variables, calculate gradients, update them, and return the new values. section 3.5.4
