How to gauge in custom JVP if a primal/tangent requires a gradient?

[atiyo]

I'm writing a custom JVP for a non-differentiable function. But the gradient can be estimated in a computationally expensive way for each argument. So I would like to restrict the gradient computation for only those components which require it. Is there a recommended way to do this?

Code example:

@custom_jvp
def f(x: list[array]):
    """a non differentiable funcion, but the gradients can be estimated in an
    computationally expensive way"""
    return g(x)


@f.defjvp
def f_jvp(primals, tangents):
    x = primals[0]
    x_dot = tangents[0]
    fwd = g(x)
    jvp = jnp.zeros_like(g(x))
    for i, xi in enumerate(x):
        # estimating dg_dxi is expensive, so i would only like to do it if xi
        # requires a gradient calculation. how should this be judged?
        if xi_grad_needed(tangents, i):
            jvp = jvp + dg_dxi_estimate(x, i) * x_dot[i]
    return fwd, jvp

I appreciate the for loop is best formatted in terms of a lax scan/foriloop: I'm just trying to illustrate with some pseudo code to get the question across clearly.

My question is: what suitable for should x_i_gradient_needed take in the above code to avoid the computation when its a gradient with respect to x[i] is not needed? Is it sufficient for it to be something like like x_dot[i] != 0.0? Or is there a better jax-esque way of doing this?

Thanks for any help!

[dfm]

I think that the symbolic_zeros argument to jax.custom_jvp.defjvp should do the trick: https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_jvp.defjvp.html

Take a look at those docs to see if you can get what you need from that!