A stateless function does not keep persistent state between calls. It receives the values it needs as inputs and returns any updated values explicitly. This is JAX’s approach to computation: unlike an operation that quietly changes a variable or a global random seed, each call can be understood from its arguments. Independent calls are easier to parallelize and distribute without synchronization. section 3.5 section 3.5.3
Random values need explicit keys
In R, successive calls to runif(3) produce different values because a hidden global seed changes. In JAX, jax$random$normal(seed_key, shape = shape(3)) receives its key explicitly. Calling it twice with the same key gives the same values. To obtain different values, create a new key with jax$random$split(seed_key, num = 1L) and pass that key to the next call. section 3.5.3
The same rule shapes a JAX training step:
- Pass trainable weights, nontrainable weights, inputs, and targets to the loss function. Use Keras model$stateless_call() to obtain outputs and updated nontrainable weights; forward passes can change those weights. section 7.4.2
- Use jax$value_and_grad(..., has_aux = TRUE) to obtain gradients while carrying the updated nontrainable weights alongside the loss. section 7.4.2
- Pass optimizer variables, gradients, and trainable weights to optimizer$stateless_apply(). Return the loss and a new state containing all three sets of variables for the next step. section 7.4.2
optimizer <- optimizer_adam()
optimizer$build(model$trainable_variables)
train_step <- function(state, inputs, targets) { #1
.[trainable_variables, non_trainable_variables, optimizer_variables] <-
state #2
.[.[loss, non_trainable_variables], grads] <- grad_fn( #3
trainable_variables, non_trainable_variables,
inputs, targets
)
.[trainable_variables, optimizer_variables] <- optimizer$stateless_apply( #4
optimizer_variables, grads, trainable_variables
)
new_state <- list( #5
trainable_variables,
non_trainable_variables,
optimizer_variables
)
list(loss, new_state)
}
The training-step listing shows the pattern end to end: unpack state, compute loss and gradients, apply updates, and return new state. section 7.4.2
Explicit state takes more work to pass around, especially in training code, but enables compilation and automatic performance optimizations. For random operations, remember that reusing a key repeats the result; for training, remember to carry forward both model and optimizer updates. section 3.5.3 section 7.4.2
