jax.value_and_grad creates a function that returns a loss value and its gradient together, avoiding a second loss computation. It mirrors jax.grad but returns a tuple: the value first, then gradients. section 3.5.3 section 3.5.1
Differentiate a squared value
Define `compute_loss(input_var)` to return `jnp.square(input_var)`, then wrap it with `grad_fn = jax.value_and_grad(compute_loss)`. Calling `output, gradient = grad_fn(input_var)` produces both the computed value and its gradient with respect to the argument. section 3.5.3
Use the same wrapper as the loss function grows:
- Pass model state as the first argument when differentiating several variables; tuples, lists, dictionaries, and nested combinations are supported as JAX trees. section 3.5.3
- Set `has_aux=True` when the loss function returns `(loss, output)`; the call then returns the loss, gradients, and auxiliary output. section 3.5.3
- Keep training functions stateless by passing parameters in and returning updated parameters alongside the loss. section 3.5.4
Linear-classifier training step
Package prediction and mean squared error in `compute_loss(state, inputs, targets)`, where `state` contains `W` and `b`. Wrap it with `jax.value_and_grad`, unpack the gradients, subtract each gradient multiplied by `learning_rate`, and return `loss, W, b`. Repeating this stateless step trains the classifier. section 3.5.4
learning_rate = 0.1 @jax.jit #1 def training_step(inputs, targets, W, b): loss, grads = grad_fn((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 return loss, W, b #4
Shows the complete gradient-based update pattern, including JIT compilation, state unpacking, parameter updates, and returned values. section 3.5.4
