Skip to content

← JAX Autodiff step 24 of 25

Medium Primitives

Saved Residuals in custom_vjp

Jump to the task

A custom reverse-mode derivative often needs information from the forward calculation. With jax.custom_vjp, the forward rule returns both the function result and saved values called residuals. JAX passes those residuals to the backward rule together with an incoming cotangent.

Save a pair of values

Consider the scalar function f(x) = sum(x * sin(x)), with angles in radians. Its derivative for each input element is:

df/dx[i] = sin(x[i]) + x[i] * cos(x[i])

The forward rule can compute sine once and reuse it for both the function value and the residuals:

def fwd(x):
    sin_x = jnp.sin(x)
    value = jnp.sum(x * sin_x)
    return value, (x, sin_x)

The return value is a pair: (primal_output, residuals). Residuals are one object, which can itself be a tuple, dictionary or another supported pytree. They need not always be a tuple, but (x, sin_x) is convenient here. Returning (value, x, sin_x) would have the wrong outer structure.

Apply the incoming cotangent

The backward rule receives a scalar cotangent g because f returns a scalar. It multiplies the saved derivative expression by that value:

def bwd(residuals, g):
    x, sin_x = residuals
    return (g * (sin_x + x * jnp.cos(x)),)

Return a one-element tuple because the primal function has one argument. Its array must match the input's shape. jax.grad(f)(x) uses an incoming cotangent of one. A general VJP can use another value, so do not hard-code one inside the backward rule.

Register the two rules with f.defvjp(fwd, bwd) after decorating the primal function with @jax.custom_vjp. The primal and forward rule must calculate the same function value.

Storage versus recomputation

Saving a value can avoid calculating it again, but uses residual storage. Recomputing a pure operation can also be a valid memory/compute tradeoff; it is not inherently incorrect. This exercise deliberately saves sine and computes cosine during backward. It does not benchmark that choice.

Keep transformed code pure and use explicit PRNG keys for randomness. Residuals do not make arbitrary Python side effects reliable under JAX transformations. custom_vjp defines reverse-mode behavior; direct forward-mode differentiation with jax.jvp is not supported by this decorator. Use custom_jvp when that is the rule you need to define.

Problem

Implement custom_vjp_with_extra_residual(x):

  1. Define the scalar f(x) = jnp.sum(x * jnp.sin(x)) with custom_vjp.
  2. In the forward rule, return its value and (x, sin_x), computing sine once for reuse within that rule.
  3. In the backward rule, return the one-input cotangent tuple above.
  4. Return jax.grad(f)(x).

x is a finite float32 JAX vector of shape (n,), with n >= 1, containing angles in radians. Leave it unchanged. Return a JAX vector of the same shape, not the scalar function value. For [0, 1, 2], the gradient is approximately [0, 1.3817732, 0.0770037].

Numerical output tests cannot distinguish a saved-residual implementation from the direct derivative formula. Check the rule functions and a VJP with a non-unit cotangent locally to verify the requested construction.