How do functions with multiple outputs compute gradients?

[DoTulip]

Hello, JAX team.
I have some confusion about the automatic differentiation of functions with multiple outputs. For a function with multiple outputs, how to get the gradients of all outputs to inputs at once? Now I achieve what I want based on VJP, but I am not sure if this is an efficient way. Below is the code I am currently using.

import jax
import jax.numpy as jnp
from jax import value_and_grad, grad, jvp, vjp
from numpy import float32

def fun_multi_outputs(xPhys):
    outarray = jnp.zeros(2)
    output_A = jnp.sum(xPhys)
    output_B = jnp.sum(xPhys ** 2)
    outarray = outarray.at[0].set(output_A)
    outarray = outarray.at[1].set(output_B)
    print('funA is running')
    return outarray

if __name__ == '__main__':

    case = 'A'

    if case == 'A':
        inputs = jnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=float32)

        fun_outputs, vjp_fun = vjp(fun_multi_outputs, inputs)

        v_vjp_A = jnp.array([1, 0], dtype=float32)
        output_A_grad = vjp_fun(v_vjp_A)[0]
        print("output_A's grad is: ", output_A_grad)

        v_vjp_B = jnp.array([0, 1], dtype=float32)
        output_B_grad = vjp_fun(v_vjp_B)[0]
        print("output_B's grad is: ", output_B_grad)

Is there a more direct and efficient way to implement it?

[YouJiacheng]

You can use jax.jacrev and jax.jacfwd.
jax.jacrev can be considered as vmapped vjp. It is basically the same as your code, but do it for all outputs in parallel.
jax.jacfwd can be considered as vmapped jvp.
I recommend you to read #47 (comment).
And https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html.

[DoTulip]

Thank you very much for your kind answer! I tried the method you introduced, but I found that with jacrev/jacfwd, it is not possible to get the output value of the function at the same time. If the output value of the function is calculated again, it will cause repeated calls of the sub-function. Is there a way to get the output value of the function at the same time when using jacrev/jacfwd?

[YouJiacheng]

@DoTulip
Definitely Yes! There are 2 ways to do so.
First, you can workaround with has_aux, namely:

def wrapper(fun):
    def fun_duplicate(*args, **kwargs):
        y = fun(*args, **kwargs)
        return (y, y)
    return fun_duplicate

jac, y = jax.jacfwd(wrapper(f), has_aux=True)(x)
jac, y = jax.jacrev(wrapper(f), has_aux=True)(x)

Second, you can refer to the source code of jax.jacrev and jax.jacfwd and implement the with value version by yourself.
Here is an example:

def value_and_jacfwd(fun: Callable[[jnp.ndarray], jnp.ndarray]):
    def value_and_jacfwd_fun(x: jnp.ndarray):
        assert x.ndim == 1
        pushfwd = partial(jax.jvp, fun, (x,))
        y, jac = jax.vmap(pushfwd, out_axes=(None, -1))((jnp.eye(x.shape[0], dtype=x.dtype),))
        return y, jac
    return value_and_jacfwd_fun

def value_and_jacrev(fun: Callable[..., jnp.ndarray]):
    def value_and_jacrev_fun(x):
        y, pullback = jax.vjp(fun, x)
        y: jnp.ndarray # y should be an array
        assert y.ndim == 1
        jac = jax.vmap(pullback)(jnp.eye(y.shape[0], dtype=y.dtype))[0]
        return y, jac
    return value_and_jacrev_fun

y, jac = value_and_jacfwd(f)(x)
y, jac = value_and_jacrev(f)(x)

[DoTulip]

Thank you for your help! Your suggestion helped me a lot!