Skip to content

← JAX Autodiff step 19 of 25

Medium Primitives

vmap(grad) vs grad(sum(vmap))

Jump to the task

Sum gradients or differentiate the sum

For a sum of differentiable losses, the derivative of the sum equals the sum of the derivatives. You can therefore compute a total batch gradient in two ways:

  1. Differentiate each example's loss, then sum the gradients.
  2. Sum the example losses, then differentiate the total.

The first route also gives you the individual gradients. The second returns only the total. Both routes use the same shared parameter and the same examples.

import jax
import jax.numpy as jnp

def loss(w, x):
    return (w * x) ** 2

w = 1.0
xs = jnp.array([1.0, 2.0])
grad_one = jax.grad(loss, argnums=0)
grads = jax.vmap(grad_one, in_axes=(None, 0))(w, xs)
way_a = grads.sum()  # [2.0, 8.0] sums to 10.0

losses = jax.vmap(loss, in_axes=(None, 0))
def total_loss(w):
    return losses(w, xs).sum()

way_b = jax.grad(total_loss)(w)  # 10.0

The individual losses are (w*1)**2 and (w*2)**2. Their sum is 5*w**2, whose derivative at w=1 is 10. Computing the two individual derivatives first gives 2 + 8, the same total.

Keep the reduction and parameter consistent

in_axes=(None, 0) shares w across the examples and selects entries from xs. Mapping over w instead would describe separate parameters, not the shared weight in this exercise.

Both routes must use a sum. A mean loss gives a mean gradient, which is smaller by the batch size for a nonempty batch. For a vector parameter with shape (D,), per-example gradients have shape (N, D); sum only axis 0 to retain one derivative for each parameter.

The equivalence applies to the original gradients. Clipping individual gradients before summing changes the result: clipping [2, 8] to the interval [-3, 3] gives [2, 3], whose sum is 5 rather than 10. Floating-point evaluation can also introduce small differences between the two unmodified routes, so compare them with a tolerance.

The diagram shows the first route before reduction. [] denotes a scalar shape. Two example gradients produce shape [2], then summing that array gives the scalar total.

Loading visualization…

Problem

Implement equivalent_two_ways(w, x_batch) for the single-example loss (w*x)**2. The shared weight w is a real floating-point scalar, and x_batch is a one-dimensional real floating-point JAX array of shape (N,).

Compute the total gradient with respect to w using each route above. Return a JAX array with shape (2,), containing [sum_of_example_gradients, gradient_of_total_loss] in that order. Do not average, clip or normalize either result.

For w=-1.0 and x_batch=[1.0, 2.0], both totals are -10, so return [-10.0, -10.0]. The formula 2*w*sum(x_batch**2) checks the answer. Implement both constructions to practice their order of operations; numerical tests alone cannot prove that you used two different routes.