value_and_grad of pytrees

[smartalecH]

Typically, we take the grad of a scalar-valued function. I know we can get clever with vmap, and efficiently take the grad of vector-valued functions. Is there a way to take a grad of a pytree? This is mostly just syntactic sugar, but I can think of a few use cases where you are performing an expensive computation which simultaneously produces multiple outputs, and you'd like to backpropagate through all of them (without manually splitting your code and forcing a redundant "forward pass" for each output variable).

Here's a simple example which complains about the Dict:

def simple_losses(x):
    return {
        "a" : jnp.sum(x),
        "b" : jnp.sum(x ** 2),
        "c" : jnp.sum(x ** 3),
    }

x = 2 * jnp.ones((5,))

val_and_grad = jax.value_and_grad(simple_loss)

print(val_and_grad(x))

[jakevdp]

You could do this directly with jax.jacrev:

grad_only = jax.jacrev(simple_loss)
print(grad_only(x))

There is no equivalent function that computes the value and jacobian in a single pass, primarily because it wouldn't save much computation in practice (see the discussion in #762)

[smartalecH]

Ah, good call! Thanks, @jakevdp

[mfschubert]

Here you could also do the following:

def simple_loss_with_aux(x):
    y = simple_loss(x)
    return y, y

jac, value = jax.jacrev(simple_loss_with_aux, has_aux=True)(x)

[smartalecH]

Nice! Thanks!