jax.grad with multiple return values

[halfroad]

Hi Folks,

Anybody knows the underlying logics when differentiating a function with multiple return values? I am trying to differentiate the function following code lines,

import jax

def function(x, y):

    # f(x, y) = 2x² * 3y³
    # g(x, y) = 4x³ + 5y?~A?
    return 2 * x ** 2 * 3 * y ** 3, 4 * x ** 3 + 5 * y ** 4

def test():

    x = 2.
    y = 3.

    # Partial derivatives = [(dx, dy), (dx, dy)] = [(4x * 3y³, 2x² * 9y²), (12x², 20y³)] = [(648,,
 648), (48, 540)]
    function_grad = jax.grad(function, **argnums = (0, 1)**, has_aux = True)
    derivatives = function_grad(x, y)

    print("derivatives = ", derivatives)

if __name__ == "__main__":

    test()

The print of execution as follows,

derivatives =  ((Array(648., dtype=float32, weak_type=True), Array(648., dtype=float32, weak_type=True)), Array(437., dtype=float32, weak_type=True))

Supposedly, the print should appear somethings like (648,,
648), (48, 540), but not. I speculate that 437 is the value of function f(x, y) = 4 * x ** 3 + 5 * y ** 4.

In the case I like to compute the partial derivative of (48, 540) for f(x, y) = 4 * x ** 3 + 5 * y ** 4 instead, how to designate the argnums = (0, 1), or other changes?

Thank you in advance.

[jakevdp]

Thanks for the question! There are two issues here with how you're using jax.grad.

1. jax.grad only works with single-output functions. Passing has_aux allows you to use a function which returns a second auxillary argument which will not be differentiated (see docs). If you want to compute the derivative of a function with multiple outputs in one pass, one way is to use a more general gradient transform such as jax.jacobian.
2. Passing argnums=(0, 1) computes df/d{x, y}, which is useful in some cases but does not appear to be what you expected. It sounds like you want to compute (df/dx, df/dy), which requires two gradient evaluations.

With these two pieces in mind, I'd probably do something like the following to compute the results you're were looking for:

df1_dx, df2_dx = jax.jacobian(function, argnums=0)(x, y)
df1_dy, df2_dy = jax.jacobian(function, argnums=1)(x, y)
result = ((df1_dx, df1_dy), (df2_dx, df2_dy))

[halfroad]

Nice, the result is what I expected. Thank you very much @jakevdp