Calculating Jacobian with respect to certain outputs only

[Justin-Tan]

Hi, is it possible to calculate the Jacobian of a function with respect to certain outputs only?

e.g. for

def my_func(x,y,z):
    return x**2,y**2,z**2

I'd like to find the Jacobian of the first output only, with respect to the function's input arguments.

[jakevdp]

I think the recommended way to do this would be something like

my_func_jac = jax.jacobian(lambda x, y, z: my_func(x, y, z)[0])

If you jit-compile the function, the XLA compiler will do dead code elimination to ensure that unused results are not computed.

[Justin-Tan]

Thanks! Would using a functools.partial call over the lambda achieve the same thing?

[jakevdp]

I'm not sure what you mean