Jax equivalent to PyTorch "needs_input_grad" for custom vjp/gradient?

[amifalk]

Computing all of the partial derivatives in a custom gradient when only some are needed can be a costly expense. PyTorch allows users to optimize for this by passing in a boolean tuple called "needs_input_grad" (True, False, ...) denoting whether or not each parameter needs its partial derivative computed. Is there a way to support this in Jax?

[anh-tong]

I think you may need to specify argnums in jax.grad (refer to the documentation here).

Let's say that you have a function $f(x, y, z) = x + y + z$

def f(x, y, z):
      return x + y + z

Computing gradient w.r.t x and z (i.e., $\frac{\partial f}{\partial x}$ and $\frac{\partial f}{\partial z}$) can be done by

# `0` is the argument position of `x`
# `2` is the argument position of `z`
jax.grad(f, argnums=(0, 2))(x, y, z)
# this a tuple containing two gradients

[amifalk]

Thanks for the response. For custom gradients, (e.g. functions decorated with @custom_jvp), argnums only selects the number of gradients to return - there's no way to avoid computation for a partial derivative that isn't requested. For reference here's a comparison between JAX and PyTorch.

# JAX backward pass with @custom_jvp 

def f_bwd(res, g):
   x, y, z = res
   dx, dy, dz = None

   # no flags to avoid computation
   # do work on dx, dy, dz ...    
    
   return (dx, dy, dz)

f.defvjp(f_fwd, f_bwd)

# PyTorch backwards pass by subclassing Function

@staticmethod
def backward(ctx, grad_output):
   x, y, z= ctx.saved_tensors
   dx = dy = dz = None

   if ctx.needs_input_grad[0]:
      # do work on dx
   if ctx.needs_input_grad[1]:
      # do work on dy
   if ctx.needs_input_grad[2]:
      # do work on dz
   return dx, dy, dz

[anh-tong]

Oh, in this case, you can do something like this

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.custom_vjp, nondiff_argnums=(3,))
def f(x, y, z, needs_input_grad=(True, True, True)):
    return x + y + z


def f_fwd(x, y, z, needs_input_grad=(True, True, True)):
    res = (x, y, z)
    return x + y + z, res

def f_bwd(needs_input_grad, res, g):    
    x, y, z = res
    dx = dy = dz = None
    if needs_input_grad[0]:
        dx = jnp.ones_like(x)
    if needs_input_grad[1]:
        dy = jnp.ones_like(y)
    if needs_input_grad[2]:
        dz = jnp.ones_like(z)
    
    return dx, dy, dz

f.defvjp(f_fwd, f_bwd)


x = y = z = jnp.array(1.)

print(jax.grad(f, argnums=(0, 1, 2))(x, y, z)) # return (1., 1., 1.)

print(jax.grad(f, argnums=(0, 1, 2))(x, y, z, (True, False, True)))  # return (1., 0., 1.)

You can check out more about the practice using nondiff_argnums in https://jax.readthedocs.io/en/latest/jep/4008-custom-vjp-update.html

[amifalk]

Thanks! I also found #1983 informative.