Question about pure callbacks, custom_jvp & transposition

[SNMS95]

I am trying to write a code which calls an external library to solve a linear system.
However, I want to both JIT compile and perform higher order auto-diff on it.
I was thinking of using pure_callback functionality to call the external library and
then attach a custom_jvp rule to achieve what I wanted.

A minimal reproducible example looks as follows:

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np

def _solve_with_numpy(A, b):
    """Mocking external library function"""
    res = np.linalg.solve(A, b)
    return res.astype(b.dtype).reshape(b.shape)

@jax.custom_jvp
def solve(A, b):
    """Jax suitable wrapper for external library function"""
    result_shape_dtypes = jax.ShapeDtypeStruct(
            jnp.broadcast_shapes(b.shape), b.dtype)
    return jax.pure_callback(
            _solve_with_numpy,
            result_shape_dtypes,
            A, b)

@solve.defjvp
def solve_jvp(primals, tangents):
    """Mock custom rule."""
    A, b = primals
    _, b_dot = tangents
    res = solve(A, b)
    res_dot = solve(A, b_dot*2.0)  # Arbitrary rule
    return res, res_dot

A = jnp.array([[1., 2.], [3., 4.]])
b = jnp.array([1., 2.])

#jax.jacfwd(jax.jit(lambda x: jnp.sum(solve(A, x))))(b) # THIS WORKS
jax.grad(jax.jit(lambda x: jnp.sum(solve(A, x))))(b). # THIS FAILS

I see that the issue is with res_dot = solve(A, b_dot*1.0) (which is an arbitrary rule) but I want to
solve the same system but with the RHS being the accumulated tangents in b.

Is there any way to make this work?

[jakevdp]

Hi - the issue is that custom_jvp only works with functions that are linear – since your gradient calls back to the original function, it is not linear, and you'll need to define a custom_vjp. For an example of how to define this, you can refer to https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#custom-vjps-with-jax-custom-vjp

[jakevdp]

If you remove the pure_callback, then you're attempting to call NumPy code with a traced value, which will lead to a TracerArrayConversionError.

Regarding linearity: it's true you're solving a linear system, but in terms of autodiff, if the tangent of a function f requires calling the function f, the function is not linear in an autodiff sense.

The solution here is to use custom_vjp instead of custom_jvp: if you do that, your problem will be fixed. I don't know what you mean when you say you want "access to both modes".

[SNMS95]

Hey Jake, thanks for the patience :) .

1. By both modes, I meant accessing forward mode as well as reverse mode AD. A custom_vjp does not allow the use of the forward mode right?
2. If the JVP rule for linear system solving is non-linear (from an AD perspective), how is it currently implemented? I know that we can't provide both a custom_vjp and custom_jvp at the same time.
   (Defining both custom_vjp and custom_jvp #3996)

[jakevdp]

custom_vjp allows for both forward and reverse mode autodiff. Edit: it doesn't actually!

[jakevdp]

Regarding your second question: JAX's built-in solvers are implemented as primitives with forward and reverse autodiff rules. They do not use custom_jvp or custom_vjp.

[SNMS95]

I found a hacky way to at least make 2nd order reverse-mode work with pure_callback.
I do not know whether this is bad practice but the following code works in reverse-mode

I define a non-sensical JVP rule for the pure_callback solve

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np

def _solve_with_numpy(A, b):
    """Mocking external library function"""
    res = np.linalg.solve(A, b)
    return res.astype(b.dtype).reshape(b.shape)

@jax.custom_jvp
def solve(A, b):
    """Jax suitable wrapper for external library function"""
    result_shape_dtypes = jax.ShapeDtypeStruct(
            jnp.broadcast_shapes(b.shape), b.dtype)
    return jax.pure_callback(
            _solve_with_numpy,
            result_shape_dtypes,
            A, b)

@solve.defjvp
def solve_jvp(primals, tangents):
    """Mock custom rule."""
    A, b = primals
    _, b_dot = tangents
    return b*123.0, b_dot*457.0

Then I wrap this with a custom_vjp rule

@jax.custom_vjp
def inner_loss_fn(b):
    x = solve(A, b)
    return jnp.sum(x)

def inner_loss_fn_fwd(b):
    x = solve(A, b)
    return jnp.sum(x), x

def inner_loss_fn_bwd(res, g):
    ones_vec = jnp.ones_like(res)
    adjoint = jnp.linalg.solve(A.T, ones_vec*g)
    return adjoint,

inner_loss_fn.defvjp(inner_loss_fn_fwd, inner_loss_fn_bwd)    

And second-order derivatives now work well

A = jnp.array([[1., 2.], [3., 4.]])
b = jnp.array([1., 2.])
jax.grad(lambda x: jnp.sum(jax.grad(inner_loss_fn)(x)))(b)

Without the fake-jvp rule, this does not work and Errors out for second-order reverse-mode derivatives!