Is it possible to specify un-traceable sections of code within auto-differentiation?

[fishjojo]

Hi,
Suppose I have the following code

def func(x):
    _ = g(x)
    return f(x)
jac = jax.jacfwd(func)(x)

You can see that g does not contribute to the gradient, and it is also un-traceable in general.
Is there a way to tell jax not to trace through g when doing auto-differentiations?
I can define custom jvp rules for g but it doesn't look elegant.

Thanks.

[shoyer]

jax.lax.stop_gradient does exactly this for automatic differentiation, e.g., _ = g(jax.lax.stop_gradient(x))

If you also want to avoid tracing with other JAX transformations, you may want to experiment with host_callback: https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html

Note: I'm presuming that you use the result of g(x) in some way, rather than just throwing it away. If not, you may run into other potential trouble, because in general JAX expects that transformations are applied to pure functions -- and functions whose return value is not used are typically used for their side-effects