Skip to content

← JAX Autodiff step 18 of 25

Medium Primitives

Per-Example Gradients via vmap(grad(...))

Jump to the task

Keep one gradient per example

A batch loss usually sums or averages individual losses. Its gradient combines their contributions too. Sometimes you need to inspect each contribution before combining them, for example to find examples with large gradients or apply a separate clipping rule to each example.

Write a scalar loss for one example, differentiate it, then use jax.vmap to apply that gradient function across the batch:

import jax
import jax.numpy as jnp

def loss(w, b, x, y):
    return (w * x - y) ** 2

grad_one = jax.grad(loss, argnums=0)
grad_batch = jax.vmap(grad_one, in_axes=(None, None, 0, 0))
grads = grad_batch(1.0, 0.0, jnp.array([1.0, 2.0]), jnp.array([0.0, 0.0]))
# grads = [2.0, 8.0]

The derivative for one example is 2*(w*x - y)*x. With w=1, the first example contributes 2 and the second contributes 8. The gradient of the summed loss would be 10; the gradient of the mean loss would be 5. Neither single number preserves the two separate contributions.

This exercise's loss deliberately ignores b, which remains in the function signature. Do not add it to the prediction.

Share parameters and map over data

in_axes=(None, None, 0, 0) keeps w and b fixed across examples. The two zeros select corresponding entries from x_batch and y_batch. Each call receives one scalar x and one scalar y. It returns one scalar gradient, so stacking N results produces shape (N,).

A vector parameter with shape (D,) would instead give per-example gradients of shape (N, D). Keeping these separate may use much more memory than computing one batch gradient. Vectorization does not remove the cost of storing the requested results.

The diagram shows the two-example calculation. [] denotes a scalar shape and [2] denotes the shape of the stacked result.

Loading visualization…

Problem

Implement per_example_grads(w, b, x_batch, y_batch) using jax.vmap over jax.grad of the single-example loss (w*x - y)**2. Differentiate with respect to w only.

w and b are real floating-point scalars. x_batch and y_batch are real floating-point JAX arrays with matching shape (N,). Return a JAX array with shape (N,), in the same order as the examples. Do not sum, average, clip or normalize the gradients. Ignore b even when it is nonzero.

For w=1.0, b=5.0, x_batch=[-1.0, 2.0], and y_batch=[1.0, 3.0], the errors are [-2, -1]. Multiplying each by twice its x gives [4.0, -4.0]. The two gradients cancel when summed, but both must appear in your result.