Need the derivative of a Python numerical function without deriving it by hand? jax.grad transforms that function into another function that computes its gradient. It fits JAX’s functional style: parameters, inputs, and other state remain explicit function arguments rather than hidden inside objects. section 4.2.3 section 1.1.2
Use the transformed function according to the shape and arguments of the computation:
- Create a derivative function with grads_fn = jax.grad(loss_fn), then evaluate it at model parameters and data; the default differentiation target is the first function argument. section 2.6.2 section 4.2.3
- Differentiate another argument or several arguments with argnums, such as jax.grad(dist, argnums=1) or argnums=(1, 2); the result preserves the corresponding gradient structure. section 4.2.3
- Differentiate functions whose parameters are nested dictionaries or other pytrees; the returned gradients use a matching nested structure. section 4.2.3 section 2.11
- Use has_aux=True when the function returns a primary scalar result plus auxiliary data, and use value_and_grad when the function value and gradients are both needed. section 4.2.3 section 4.4
- Compose grad with itself for higher-order derivatives, with vmap for per-sample or vectorized evaluation, and with jit for compiled execution. section 4.2.4 section 6.3.3
Gradient update for a linear model
Define a pure loss function whose first argument contains the model parameters, transform it, evaluate the transformed function, and update the parameters: model_parameters = jnp.array([1., 1.]); grads_fn = jax.grad(loss_fn); grads = grads_fn(model_parameters, xt, yt); model_parameters -= learning_rate * grads. Unlike PyTorch’s requires_grad and backward workflow, this approach marks no tensors and exposes the gradient computation as an ordinary function. section 4.2.3 section 4.2.2
Higher-order differentiation
For f(x) = x⁴ + 12x + 1/x, define f_d1 = jax.grad(f), f_d2 = jax.grad(f_d1), and f_d3 = jax.grad(f_d2). Each result is another callable derivative function that can be evaluated at a chosen point. section 4.2.4
Listing 4.9 Calculating gradients in JAX
import jax #1
import jax.numpy as jnp #1
xt = jnp.array(x) #2
yt = jnp.array(y) #2
learning_rate = 1e-2
model_parameters = jnp.array([1., 1.]) #3
def model(theta, x): #4
w, b = theta
return w * x + b
def loss_fn(model_parameters, x, y): #5
prediction = model(model_parameters, x)
return jnp.mean((prediction-y)**2)
grads_fn = jax.grad(loss_fn) #6
grads = grads_fn(model_parameters, xt, yt) #7
model_parameters -= learning_rate * grads #8
#1 Imports PyTorch
#2 Converts the training data to JAX arrays
#3 Tensors with model weights without any special mark
#4 The function implements a simple linear model.
#5 MSE loss function
#6 Creates a function for calculating gradients
#7 Calculates gradients
#8 Makes a single step of gradient update
Shows the minimal autodiff pattern: create df = jax.grad(f), evaluate df(x), and inspect the derivative value. section 4.1.4
