Evaluate jax.grad(f) with respect to particular leaves in a PyTree

[conorheins]

This question seems very similar to this one but I wanted to write it again here to both increase visibility of the issue, as well as to make the question both more concise/clear and also more general.

Consider a function f with scalar output that takes a PyTree as input:

def f(input: PyTree): -> float
        loss = do_stuff(input)
        return loss

Is there a quick/JAX-ian way to compute jax.grad(f) with respect to only certain leaves of the input PyTree, rather than with respect to all the leaves (the default).

My naïve approach would be to map lax.stop_gradient over all the leaves of the PyTree that you are not interested in getting gradients with respect to, leaving only the leaf-of-interest to have a non-zero gradient evaluation. But I'm not sure that's the most efficient/JAXian way to do this.

Thanks in advance!

[jakevdp]

No, there's no mechanism to do this aside from defining a function in terms of the pytree elements that you're interested in differentiating against. Options for this have been proposed (e.g. #3875 and #10614) but nothing has yet been implemented. Your best bet is probably to compute the gradient with respect to flattened arguments; it's a bit messy, but the workaround might look something like this:

from jax import tree_util, grad
import jax.numpy as jnp
from typing import NamedTuple

class Data(NamedTuple):
  x: jnp.ndarray
  i: int

data = Data(jnp.arange(5.0), 1)

# We want to compute the gradient of f with respect to x... how?
def f(data):
  return data.x[data.i]

# Wrap it with a flattened version of the arguments...
def f_wrap(x, i, tree):
  data = tree_util.tree_unflatten(tree, (x, i))
  return f(data)

args, tree = tree_util.tree_flatten(data)
result = grad(f_wrap)(*args, tree=tree)

[conorheins]

Thanks very much for the quick answer @jakevdp -- much appreciated & good to know!