How do you nest implementations of VJP?

If I have a function f that has a VJP. Can I somehow build a VJP on top of it for some function g? My use case is that the backwards pass of g is quite complicated and makes multiple calls to f_vjp. Here’s a simplified example:

from jax import vjp, custom_vjp

def f(x):
    return x ** 2

@custom_vjp
def g(x):
    return f(x)

def g_fwd(x):
    return vjp(f, x)

def g_bwd(f_vjp, y_bar):
    return f_vjp(y_bar)

g.defvjp(g_fwd, g_bwd)
y, g_vjp = vjp(g, 1.0)

This prints: TypeError: <class 'functools.partial'> is not a valid JAX type because vjp returns a functool.partial instance instead of a pytree. It should be possible for vjp to return a pytree-like callable instead since internally vjp(f, x) produces residuals that must be pytree-like. It’s just unfortunate that when they’re wrapped up into a callable, then that callable is not pytree-like.

[shoyer]

This seems like a case where perhaps you are pushing the limits of what custom_vjp is designed to support should be considering other options 😃

That said, one way to fix this immediate issue would be to replace partial inside jax.api._vjp with tree_util.Partial, which is serializable as a pytree. You could also do this wrapping in user code:

from jax import vjp, custom_vjp, tree_util

def f(x):
    return x ** 2

@custom_vjp
def g(x):
    return f(x)

def g_fwd(x):
    y, f_vjp = vjp(f, x)
    return y, tree_util.Partial(f_vjp)

def g_bwd(f_vjp, y_bar):
    return f_vjp(y_bar)

g.defvjp(g_fwd, g_bwd)
y, g_vjp = vjp(g, 1.0)
print(y, g_vjp(1.0))
# 1.0 (DeviceArray(2., dtype=float32),)

In general, it’s fine to using Partial inside the forward pass of custom_vjp functions as long as you are careful not to close over any tracers, like the value x in this case. The dtypes, constants, jaxprs and treedefs used in the closures should all be fine. (Otherwise you should get a nasty error message.)

(@mattjj please correct me if I’m mis-stating anything here)

[NeilGirdhar]

Thanks for your help. I made a pull request that addresses this issue. When you have time, do you mind taking a look?