jax.grad transforms a loss function into another function that computes gradients. Define the loss, create the gradient function, then call it with the original arguments. section 3.5.3
Differentiate a scalar loss
Define a loss that returns a scalar, transform it with jax.grad, and evaluate the result: compute_loss(input_var) returns jnp.square(input_var); grad_fn = jax.grad(compute_loss); grad_fn(jnp.array(3.0)) returns the loss gradient with respect to input_var. section 3.5.3
Use these rules when applying jax.grad:
- Return a scalar loss from the differentiated function. section 3.5.3
- Place the state arrays requiring gradients in the first argument; that argument may be one array, a list, or a dictionary of arrays. section 3.5.3
- Use jax.value_and_grad when the computation must return both the loss and its gradients, including optional auxiliary outputs. section 3.5.1 section 3.5.3
Listing 10.13 Loss maximization via stochastic gradient ascent: JAX
import jax grad_fn = jax.grad(compute_loss) @jax.jit def gradient_ascent_step(image, filter_index, learning_rate): grads = grad_fn(image, filter_index) grads = ops.normalize(grads) image += learning_rate * grads return image
A JAX gradient function can drive a JIT-compiled gradient-ascent step that normalizes gradients and updates an input image. section 10.2.3
