Tracing into function with specific custom_jvp

[SNMS95]

I am trying to create a custom jvp rule for a newton ralphson root finder algorithm.
Since I use linear solvers and host of other external stuff inside the newton_solver_fn(), I do the following:

class NewtonSolver(equinox.Module):
   def  _solve_with_backend(...):
        # While loop from lax
              # pure_callback to external library to solve (Ax=b)
        return x_final
              
 
   def solve(...): # User API
        sol = wrapped_solve(self._solve, ...)
        return sol

# ============

@partial(custom_jvp, nondiff_argnums=(0,))
def wrapped_solve(solve_fn, ...):
     sol = solve_fn(...)
     return sol
     
@wrapped_solve_fn_for_fwd_mode.defjvp
def wrapped_solve_fn_for_fwd_jvp(solve_fn, primals, tangents):     
      # a. Call wrapped_solve() -> solves Ax=b -> primal_output
      # b. Create a JVP for the "interested fn (whose root is to be found)" at the optimal point
      # c. Call wrapped_solve() again  -> Solves Ax=Jvp -> Tangent_output

As per my understanding, JAX should not look inside "_solve()".
But when I try to do jax.grad(), it shows an error:

"ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values"

which I believe is coming from inside the "_solve()".

Why is this happening? Is there something that I am missing?

P.s. The error only happens when I create the jvp [which is an abstract tracer]and use that as the RHS. If I simply do it with some other RHS vector, the code works [Of course, the gradient is wrong].

[jakevdp]

Can you put together a reproducible example of what you're seeing? custom_jvp should work with while_loop in most cases, but the details are important. It may be that you'll need to use custom_vjp to define custom reverse-mode autodiff.

[SNMS95]

Hey @jakevdp

Sorry for the late reply. There was a lot of custom code that had to be stripped down. I even took down the while loop (so that I end up with just a pure_callback to a linear solver).

import jax.numpy as jnp
import numpy as onp
import jax
from jax import custom_jvp
from functools import partial


def _solve(A_fn, params):
    """A blackbox solve function from numpy"""
    A, b = A_fn(params)
    # create a pure_callback function to solve the linear system
    _np_call_back = lambda A, b: onp.linalg.solve(A, b).astype(b.dtype)
    # Define the expected shape & dtype of output.
    result_shape_dtype = jax.ShapeDtypeStruct(shape=b.shape, dtype=b.dtype)
    sol = jax.pure_callback(_np_call_back, result_shape_dtype, A, b)
    return sol, 0.0

# Create a wrapper around it
@partial(custom_jvp, nondiff_argnums=(0, 1))
def wrapped_solve(_solve, A_fn, params):
    return _solve(A_fn, params)

# Custom JVP rule
@wrapped_solve.defjvp
def _solve_jvp(_solve, A_fn, primals, tangents):
    (params,) = primals
    (params_dot,) = tangents
    # A. primals output
    sol, err = wrapped_solve(_solve, A_fn, params)
    print("First call done")

    # B. tangent output
    new_b = params_dot * sol
    sol_dot, _ = wrapped_solve(_solve, A_fn, new_b)
    return (sol, err), (sol_dot, 0.0)


# Test the wrapped_solve function
def loss_fn(rho):

    def get_A_and_b(params):
        A = jnp.array([[3.0, 2.0], [1.0, 4.0]]) * params
        b = jnp.array([5.0, 6.0]) * (params / 2)
        return A, b

    sol, _ = wrapped_solve(_solve, get_A_and_b, rho)
    return jnp.mean(sol)

print(jax.grad(loss_fn)(1.0))

[SNMS95]

Now the error is changed to "ValueError: Pure callbacks do not support transpose. Please use jax.custom_vjp to use callbacks while taking gradients." However, I suspect that the issue is the same: tracing the blackbox "_solve()"

[SNMS95]

The issue goes away if I dont use new_b = params_dot * sol

[SNMS95]

@jakevdp Is this expected behavior or a bug?

[jakevdp]

This is expected: if you want to use the pure callback to compute the gradient, then it will need a custom_vjp in order to be reverse-mode differentiable. The error message correctly points to this.