`jax$value_and_grad()` turns a scalar-loss function into a function that returns both the loss and its gradients. Use it when a training step needs to report the loss and update weights: unlike `jax$grad()`, it supplies the loss without a separate call to compute it. section 3.5.3
Get a value and gradient together
```r compute_loss <- function(input_var) { jnp$square(input_var) } grad_fn <- jax$value_and_grad(compute_loss) input_var <- jnp$array(3) .[loss, gradient] <- grad_fn(input_var) ``` `grad_fn()` accepts the same input as `compute_loss()`. Its first result is the loss value; its second is the gradient with respect to the first argument. section 3.5.3
For several weights, put them together in the loss function’s first argument, `state`, and pass other inputs afterward. The state can be a list or a nested tree. In a linear classifier, `state` holds `W` and `b`; the resulting gradients are used to calculate new weights, which a stateless training step returns rather than modifying in place. section 3.5.3 section 3.5.4
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 training step gets loss and gradients in one call, updates both classifier weights, and returns the loss with the new weights. It also wraps the step in `jax$jit()`. section 3.5.4
If computing the loss also produces information that must be kept—such as updated nontrainable weights—return it alongside the scalar loss and set `has_aux = TRUE`. The auxiliary result is then available separately from the gradients. Keras uses this pattern with `stateless_call()` before applying gradients through an optimizer’s `stateless_apply()` method. section 3.5.3 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 Keras training step carries trainable weights, nontrainable weights, and optimizer variables explicitly through one stateless update. section 7.4.2
