What if a loss function must return more than its scalar loss—such as predictions for logging or model state for BatchNorm? Set `has_aux=True` on `jax.grad` when the function returns `(output, auxiliary_data)`. JAX differentiates only the first value and passes the auxiliary data through. section 4.2.3
Return predictions while computing gradients
```python def loss_fn(model_parameters, x, y): prediction = model(model_parameters, x) return jnp.mean((prediction - y) ** 2), prediction grads_fn = jax.grad(loss_fn, has_aux=True) grads, preds = grads_fn(model_parameters, xt, yt) model_parameters -= learning_rate * grads ``` Without `has_aux=True`, `grad()` rejects the tuple because ordinary gradient transformation expects a scalar function result. With it enabled, `grads` contains the parameter gradients and `preds` contains the auxiliary predictions. section 4.2.3
Use `has_aux` when the differentiated computation also needs to expose useful results.
- Maintain internal model state, such as running statistics used by BatchNorm. section 4.2.3
- Combine the loss value, auxiliary data, and gradients with `jax.value_and_grad(loss_fn, has_aux=True)`, which returns `((loss, auxiliary_data), gradients)`. section 4.2.3
