Higher-Order custom_vjp

A custom VJP supplies a reverse-mode derivative rule. You can then differentiate the gradient computation again, provided its operations support the required transformations and retain the needed dependence on the inputs. This exercise uses a quadratic function whose result can be checked by hand.

First derivative

For f(x) = sum(x ** 2), the gradient is the vector 2 * x:

@jax.custom_vjp
def custom_sumsq(x):
    return jnp.sum(x ** 2)

def fwd(x):
    return jnp.sum(x ** 2), x

def bwd(x, cotangent):
    return (cotangent * 2 * x,)

custom_sumsq.defvjp(fwd, bwd)

The forward rule saves x. The backward rule multiplies by the incoming scalar cotangent and returns a one-element tuple for the one input. jax.grad(custom_sumsq)(x) uses cotangent one.

Which second derivative are we returning?

jax.grad requires a scalar output. The first gradient is a vector, so jax.grad(jax.grad(custom_sumsq))(x) is not valid for vector x. Instead, sum that gradient before applying grad again:

grad_fn = jax.grad(custom_sumsq)
result = jax.grad(lambda z: jnp.sum(grad_fn(z)))(x)

Here the scalar being differentiated is sum(2 * x), so every result entry is 2. For x = [1, 2, 3], the first gradient is [2, 4, 6] and the returned vector is [2, 2, 2].

This vector is not the full Hessian matrix. The Hessian of this function is 2 * I, with shape (n, n). Differentiating the sum of gradient entries produces the sums of the Hessian's columns. For this separable quadratic, those happen to equal its diagonal entries. For a coupled function such as (x[0] + x[1]) ** 2, the diagonal is [2, 2] but the gradient of the summed gradient is [4, 4].

Keep the rule differentiable

Use JAX operations in the rule. Converting a traced value to a Python number or branching on its value with Python control flow can fail during transformation. Static Python loops and configuration-based branches are not categorically forbidden.

Residuals may be supported pytrees containing arrays and scalars; they need not all be arrays. What matters here is preserving the saved input's differentiable dependence. Applying stop_gradient to it would change the second derivative even if the first derivative still looked correct. Likewise, a custom rule that intentionally substitutes a surrogate first derivative does not automatically recover the primal's true Hessian.

Direct forward-mode jax.jvp through a custom_vjp primal is unsupported. This exercise uses reverse-over-reverse differentiation.