How to define the jacobian of a vector-valued func with custom_vjp where the jacobian is computed on the forward pass?

[antalszava]

Hi everyone, given a function that computes both the value and the jacobian of the function on the forward pass, what could be a way to register the jacobian when using custom_vjp?

The following is an attempt at passing the jacobian to wrapped_exec_bwd, however, it fails due to output issues:

params = jnp.array([0.1, 0.2])

@jax.custom_vjp
def wrapped_exec(params):
    res = params ** 2, params ** 3
    jacs = jnp.diag(2 * params), jnp.diag(3 * params ** 2)
    return res, jacs

def wrapped_exec_fwd(params):
    res, jacs = wrapped_exec(params)
    return res, tuple([jacs, params])

def wrapped_exec_bwd(params, g):
    jacs, params = params
    return jacs # <--- how to define the return here?

wrapped_exec.defvjp(wrapped_exec_fwd, wrapped_exec_bwd)
jax.jacobian(wrapped_exec)(params)

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
<ipython-input-56-29e210b9faaa> in <module>
     16 
     17 wrapped_exec.defvjp(wrapped_exec_fwd, wrapped_exec_bwd)
---> 18 jax.jacobian(wrapped_exec)(params[0])

~/anaconda3/lib/python3.8/site-packages/jax/_src/api.py in jacfun(*args, **kwargs)
   1033     y, pullback = _vjp(f_partial, *dyn_args)
   1034     tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
-> 1035     jac = vmap(pullback)(_std_basis(y))
   1036     jac = jac[0] if isinstance(argnums, int) else jac
   1037     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args

    [... skipping hidden 8 frame]

~/anaconda3/lib/python3.8/site-packages/jax/_src/custom_derivatives.py in _flatten_bwd(in_tree, in_avals, out_trees, *args)
    553            "number of arguments to the primal function, but got VJP output "
    554            "structure {} for primal input structure {}.")
--> 555     raise TypeError(msg.format(in_tree2, in_tree)) from None
    556   yield [zeros_like_aval(aval.at_least_vspace()) if ct is zero else ct
    557          for aval, ct in zip(in_avals, cts_in_flat)]

TypeError: Custom VJP rule must produce an output with the same container (pytree) structure as the args tuple of the primal function, and in particular must produce a tuple of length equal to the number of arguments to the primal function, but got VJP output structure PyTreeDef(((*, *),)) for primal input structure PyTreeDef((*,)).

The traceback contains the vmap call, although not sure how to make the vectorization be compatible with the computed jacobian.

[YouJiacheng]

params = jnp.array([0.1, 0.2])

@jax.custom_vjp
def wrapped_exec(params):
    y = params ** 2, params ** 3
    # don't need compute jacs here
    return y

def wrapped_exec_fwd(params):
    y = wrapped_exec(params)
    jacs = jnp.diag(2 * params), jnp.diag(3 * params ** 2) # compute here
    return y, jacs # don't need params here

def wrapped_exec_bwd(res, g):
    jac1, jac2 = res
    g1, g2 = g
    return (g1 @ jac1) + (g2 @ jac2),

wrapped_exec.defvjp(wrapped_exec_fwd, wrapped_exec_bwd)
jax.jacobian(wrapped_exec)(params)

[antalszava]

Thank you!