How to not compute discarded VJPs in the backward propagation

[fishjojo]

Dear developers,
I have a function which takes two arguments, e.g.,

@jax.custom_vjp
def f(x, y):
    return x**2 + y**2

def f_fwd(x, y):
    return f(x, y), (x,y)

def f_bwd(res, g):
    x, y = res
    return g * 2 * x, g * 2 * y

f.defvjp(f_fwd, f_bwd)

But if I only need to compute the derivative with respect to one of the arguments, is there any way to tell the backward pass not to compute the VJP for the other argument? I encounter this problem in implicit differentiation, where one of the arguments may be closed over, yet both VJPs are computed and it introduces a lot of overhead. It would be nice to know which VJP is needed before the backward pass, e.g., after the forward pass, I get a flag and then maybe I can define the backward pass like the following

def f_bwd(flag, res, g):
    x, y = res
    if flag == (0,1):
        return g * 2 * x, g * 2 * y
    elif flag == (0,):
        return g * 2 * x, None
    else:
        return None, g * 2 * y

Thanks

[amifalk]

You might find this response helpful. You'll likely want to pass in a tuple containing information about the gradients you wish to compute as a non-differentiable argument.

[fishjojo]

Thanks a lot, this helps.

[jakevdp]

One thing to keep in mind is that the XLA compiler does automatic dead code elimination, so as long as you wrap your computation in jit, there's no need to worry about doing unnecessary computations: if the computation is unused in your program, the compiler will remove it.