Skip to content

← JAX Autodiff step 25 of 25

Medium Primitives

Grad through stop_gradient

Jump to the task

Hold one input constant for differentiation

Suppose a model multiplies a feature vector x by a trainable scalar w. Its loss is sum((x*w)**2). We want the features to affect the loss and the update to w, while blocking derivatives through the feature computation itself. Apply lax.stop_gradient to x before multiplying.

This keeps x's values. It changes how automatic differentiation treats that path, not the numerical loss. With x=[1, 2] and w=3, the loss is 45 and the derivative with respect to w is 30:

import jax
import jax.numpy as jnp
from jax import lax

def loss(x, w):
    return jnp.sum((lax.stop_gradient(x) * w) ** 2)

x = jnp.array([1.0, 2.0])
loss(x, 3.0)  # 45
dx, dw = jax.grad(loss, argnums=(0, 1))(x, 3.0)
# dx = [0, 0]
# dw = 30

Without the stop, the loss and dw are unchanged, but dx becomes [18, 36]. Stopping w instead would block dw while leaving a derivative through x. Stop only the path you intend to exclude.

Constant for autodiff does not mean a fixed value

A different x can still produce a different loss and a different dw. For x=[2, 2] and w=3, dw is 48. Even so, differentiating dw with respect to x through this stopped computation returns zeros. The stop applies to nested differentiation as well as first derivatives, and to both forward- and reverse-mode autodiff.

Stopping a feature path does not freeze every parameter update in a training system. Other paths, weight decay or existing optimizer momentum can still change a parameter. The optimizer must also be configured to leave frozen parameters alone.

Problem

Implement grad_through_stop(x, w). The input x is a one-dimensional real floating-point JAX array and w is a real floating-point scalar. Define the scalar loss sum((lax.stop_gradient(x)*w)**2) and use jax.grad to return its derivative with respect to w only:

result = 2 * w * sum(x**2)

Return the derivative, not the loss or a gradient with respect to x. For x=[1, 2, 3] and w=1, return 28. For x=[-1, 2] and w=-2, return -20. The sum reduction is essential: grad requires a scalar output, whereas the unreduced squared array has one value per feature.

The output-only tests cannot detect whether you omitted stop_gradient, because the w derivative has the same value either way. Use the joint-gradient example above to check the stopped x path too.