Use jax.value_and_grad when a training step needs both the loss value and its gradient. It evaluates the loss once, avoiding a second loss calculation just to monitor learning progress. section 4.2.3 section 2.11
A training update that reports loss
from jax import value_and_grad def update(params, x, y, epoch_number): loss_value, grads = value_and_grad(loss)(params, x, y) lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS) new_params = [ (w - lr * dw, b - lr * db) for (w, b), (dw, db) in zip(params, grads) ] return new_params, loss_value The result is a pair: the scalar loss value and gradients matching the parameter structure. Returning the loss makes progress tracking part of the update rather than a separate computation. section 2.6.3
Extend the return structure when the loss function provides auxiliary information:
- Define loss_fn to return (loss, predictions), then call value_and_grad(loss_fn, has_aux=True). The result has the form ((loss, predictions), grads). section 4.2.3
- Pass parameters, inputs, and targets explicitly; JAX commonly keeps model parameters in pytrees rather than object state. section 2.11
- Use value_and_grad for first-order gradients; use grad repeatedly for higher-order derivatives and related transformations for Jacobians or Hessians. section 4.4
Listing 2.12 Implementing the gradient update step with both loss value and gradient
from jax import value_and_grad #1
def update(params, x, y, epoch_number):
loss_value, grads = value_and_grad(loss)(params, x, y) #2
lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS)
return [(w - lr * dw, b - lr * db)
for (w, b), (dw, db) in
zip(params, grads)], loss_value #3
#1 Importing value_and_grad function
#2 Calculating both the loss value and the gradient
#3 Returning both updated parameters and the loss value
Inspect the complete gradient-update listing, including the combined loss-and-gradient return. section 2.6.3
