Does JAX automatically adjust the tangent output in custom_jvp?

[ToshiyukiBandai]

Hello everyone,

I am trying to write a custom Jacobian-vector product rule for the Newton method. At first, I used custom_vjp, but I switched to custom_jvp to use both forward and reverse automatic differentiation.

When implementing custom_jvp, I encountered an interesting behavior of custom_jvp. In the example below, I want to compute the derivative $\frac{\partial \mathbf{x}}{\partial \alpha}$ and $\frac{\partial \mathbf{x}}{\partial \beta}$ by implicit differentiation given
$\mathbf{F}(\mathbf{x}, \alpha, \beta) = \mathbf{0},$
where $\mathbf{x} = (x_1, x_2)$ and $\mathbf{F}_1 = x_1^2 + x_2^2 - \alpha$ and $\mathbf{F}_2 = \beta x_1^3 - x_2 - \alpha$

from functools import partial
import jax
import jax.numpy as jnp
from jax import grad, jacfwd, hessian, jacrev, custom_vjp, custom_jvp, jvp
from jax.lax import while_loop
from jax.config import config
config.update("jax_enable_x64", True)

def f(x, a, b):
    # initialize residual
    F  = jnp.zeros(2)
    
    # compute each component in the residual
    F_0 = x[0]**2 + x[1]**2 - a
    F_1 = b*x[0]**3 - x[1]
    
    # insert the values into the residual
    # Because JAX array (F) is not mutable, we need to use a method .at[].set().
    F = F.at[0].set(F_0)
    F = F.at[1].set(F_1)
    
    return F

Jacobian_fun = jacfwd(f, argnums= 0, has_aux=False)

@partial(jax.custom_jvp, nondiff_argnums=(0,))
def newton_solver(fun, x_guess, *args):
    def cond_fun(carry):
        x_k, residual = carry
        # update tolerance
        error = jnp.linalg.norm(residual, ord = jnp.inf)
        
        return error > 1.0e-8
        
    def body_fun(carry):
        x_k, residual = carry
        J = Jacobian_fun(x_k, *args)
        # solve the linear system
        dx = jnp.linalg.solve(J, -residual)
        
        x_k = x_k + dx
        residual = fun(x_k, *args)

        return x_k, residual
    
    init_carry = (x_guess, fun(x_guess, *args))
    x_star, _ = while_loop(cond_fun, body_fun, init_carry)
    return x_star

@newton_solver.defjvp
def newton_solver_jvp(fun, primals, tangents):
    x_guess, args = primals[0], primals[1:]
    tangents_x, tangents_args = tangents[0], tangents[1:]
    primal_out = f(x_guess, *args)
    _, u = jvp(partial(fun, primal_out), args, tangents_args, has_aux=False)
    J = Jacobian_fun(primal_out, *args)
    tangent_out = jnp.linalg.solve(J, -u)
    return primal_out, tangent_out

x_guess = jnp.asarray([1.0, 1.0])

a = 1.0
b = 6.0

args = (a, b)

derivative_newton = jacrev(newton_solver, argnums=(1,2,3), has_aux=False)

derivative_newton(f, x_guess, *args)

Then, the output was

(Array([[0., 0.],
        [0., 0.]], dtype=float64),
 Array([0.00549451, 0.0989011 ], dtype=float64, weak_type=True),
 Array([-0.05494505,  0.01098901], dtype=float64, weak_type=True))

This result was unexpected to me. I expected that there was an error because tangent_out in newton_solver_jvp is not supposed to contain the derivative $\frac{\partial \mathbf{x_{solution}}}{\partial \mathbf{x_{guess}}}$ but JAX somehow filled this with the zero array, which in fact I wanted. When I implemented custom_vjp, I explicitly told JAX to fill the derivative with the zero array by putting a None tuple to the output of the reverse mode. Could anyone explain what is going on here?

[jakevdp]

Thanks for the question!

When you define a custom_jvp, you are implicitly defining the vjp as well via automatic transposition of the custom jvp function. Because the output of newton_solver_jvp has no dependence on x_guess, the automatic transposition makes this effectively equivalent to the VJP rule returning None, and so the autodiff machinery returns a zero gradient in both cases.

Does that make sense?

[ToshiyukiBandai]

Thank you for your quick response! It makes sense and, I think this is a good design! Thank you!

[patrick-kidger]

This might be coming in too late to be useful -- but heads-up that we now have a Newton implementation available in Optimistix, which might save you some implementation time. (This library is the latest entry in the Equinox scientific ecosystem, offering nonlinear optimisation.)