Couple JAX with Autograd

[RRiva]

I'm writing a package using JAX, which calls another one based on Autograd. How can I propagate the gradient?

I can think of two options:

1. Both packages compute their own gradient. Then, I will couple them together like OpenMDAO would do with different components.
2. Convert back and forth between JAX and Autograd arrays.

The second option seems better, since I won't have to write the chain rule myself, but requires knowing how dual numbers are handled by JAX and Autograd. Can you please help me?

Here is a simple code to develop this feature.

from jax import vmap
import jax.numpy as jnp
import autograd.numpy as anp
from jax import grad as jgrad
from autograd import elementwise_grad as a_elem_grad


def fn_jax(x):
    y = x**2
    z = jnp.sin(y)
    return z


def fn_autograd(x):
    y = x**2
    z = anp.sin(y)
    return z


def external_fn(x):
    return anp.sin(x)


def fn_from_jax_to_autograd(x):
    # x is a JAX array.
    y = x**2
    # TODO: convert y from JAX to Autograd array
    # TODO: define a custom derivative rule where we call AutoGrad grad on external_fn.
    z = external_fn(y)
    # TODO: convert z from Autograd to JAX array.
    return z


jx = jnp.linspace(-3, 3, 7)
ax = anp.linspace(-3, 3, 7)

res_jax = vmap(jgrad(fn_jax))(jx)
res_autograd = a_elem_grad(fn_autograd)(ax)

res_test = jgrad(fn_from_jax_to_autograd)(jx[0])
# res_test = vmap(jgrad(fn_from_jax_to_autograd))(jx)

[Dhia-naouali]

hii, I think @jax.custom_vjp is what your looking for, it basically let jax treat a certain function (external_fn as a black box when computing the gradient, by providing a forward & backward implementations that would convert and use autograd to do so

I would also recommend to take a look at torch2jax (I'm not a 100% sure, but it may solve the problem in a convenient way)

[RRiva]

Hi @Dhia-naouali, thank you for your answer 🙂 Indeed, jax.custom_vjp and jax.custom_jvp are the first step, but the difficult part is to convert back and forth between JAX and Autograd arrays. The package that you've linked is for torch.autograd, which is a different one although it has a similar name.