Gradients of functions which accept Pytrees with arbitrary leaf types (how to ignore Int)?

[femtomc]

Hi all!

I'm defining functions which accept Pytree instances -- and some of my classes have fields which are fields which I do not wish to take the gradient with respect to.

I expose these functions to users -- and in the semantics of my functions, taking the gradient with int_allow = True doesn't make sense. So I'd like to prevent it.

Is there a way to expose a wrapped function whose gradient with respect to full Pytree instances ignores non-differentiable fields?

Another solution would suffice:

Is it possible to take a function which accepts the full Pytree instance -- and wrap it so that it closes over reconstructing the instance when the differentiable parts are provided? Then maybe I could provide a sort of filter which ignores non-differentiable fields.

Potential other other solution:

If I allow allow_int = true, what if I set float0 fields in my gradients to 0 after the fact? Is this typically how people work around this problem?

The issue is that I'm doing gradient descent with Pytree instances, so later on I use a tree_map(lambda v1, v2: v1 + tau * v2, base, grad_tree).

[femtomc]

Ah, I found this: #12765

Hmm -- so, I have control over the functions I pass over to the users to take jax.grad on.

What if I define a split function which takes a Pytree, and splits it along differentiable and non-differentiable parts -- and then provide a fn(differentiable, nondifferentiable, tree)?

Has anyone tried this before?

[jakevdp]

Thanks for the question – unfortunately there's not a great solution for this at the moment. We encountered a similar problem in jax.experimental.sparse, in which we want autodiff to operate with respect to the data buffer, but not the indices buffer. Our solution there was to provide sparse-specific autodiff wrappers that know how to ignore the integer components in the sparse pytree: http://go/jax-github/blob/main/jax/experimental/sparse/ad.py

We've been talking about ways to improve on this, for example by allowing pytrees to register their preferred autodiff behavior in the same way that register_vmappable allows user-defined vmap behavior (see how it's used in e.g. http://go/jax-github/blob/main/jax/experimental/sparse/bcoo.py) @froystig and @mattjj have been thinking through that, but unfortunately nothing has landed in the JAX package yet. Until either that or the mechanism discussed in #12765 lands, for now I suspect the best course would be to mimic what we did in sparse.ad, although this is admittedly a pretty heavy-weight solution.

Thanks for the feedback - I hope we can provide a lighter-weight solution to this kind of situation in the near future.

[femtomc]

Thanks @jakevdp, no problem!

[femtomc]

@jakevdp If Pytree leaves have None value -- will that cause an issue for jax.grad? I'm not sure how jax.grad treats None.

[jakevdp]

I think None should be fine; for example here it is within a list (which is one of the pre-registered pytrees):

import jax

def f(L):
  return L[0]

print(jax.grad(f)([0.5, None]))
# [DeviceArray(1., dtype=float32, weak_type=True), None]