How to avoid tracing certain code block with jax.grad

[pipme]

N = 60
def f1(x):
    J = jnp.sum(x**2)
    J_no_grad = jnp.zeros((N, N))
    for i in range(N):
        for j in range(N):
            J_no_grad = J_no_grad.at[i, j].set(J + x + i * j)
    return J, J_no_grad

def f2(x):
    J = jnp.sum(x**2)
    return J, None

jax.grad(f1, has_aux=True) takes 5 seconds compared to 10ms with jax.grad(f2, has_aux=True). In this case I just want to get the gradient of J. J_no_grad is for recording some intermediate values. How could I reduce the computation time? Can we do something like:

with jax.no_grad():
    J_no_grad = jnp.zeros((N, N))
    for i in range(N):
        for j in range(N):
            J_no_grad = J_no_grad.at[i, j].set(J + x + i * j)

Thanks!

[pipme]

I am aware of #1937 but still no idea how to proceed in my case. There are actually a bunch of variables like J_no_grad and their computations are quite intensive, e.g. involving linalg.cholesky. Thanks for help in advance : ) @mattjj @jakevdp

[pipme]

Yes I would need jit for the function. So eventually it will be jax.jit(jax.grad(f1, has_aux=True)). When jitting it becomes even slower -- 50 seconds in the first run. The problem is that I will need to jit the function multiple times since the shape of input keeps changing. Therefore it's unacceptable for a single jit to take a long time.

[mattjj]

I would guess that the reason execution is slow in the above code without jit, even without any autodiff, is that the updates to J_no_grad are not in-place. They only become in-place updates when under a jit.

Under a jit we may have a different issue, which is that compilation times could be long if N is large. Typically if a loop is really needed this is where we'd use jax.lax.scan, but I'm not clear on whether a loop is really needed (see below).

So I suspect narrowly answering your original question may not be enough to get good performance the way this code is written. However, just to answer your question directly, you could write stop_gradient on any occurrence of x where you don't care about differentiation:

N = 60
def f1(x):
    J = jnp.sum(x**2)
    J_no_grad = jnp.zeros((N, N))
    for i in range(N):
        for j in range(N):
            J_no_grad = J_no_grad.at[i, j].set(J + jax.lax.stop_gradient(x) + i * j)  # change on this line!
    return J, J_no_grad

Could you share a fully runnable self-contained example? Is the computation of J_no_grad.at[i, j].set(J + x + i * j) actually what you care about, or is it just a toy example? (I ask the latter because it doesn't need to be expressed as a (nested) loop at all, and instead can just be written as J_no_grad = J + jax.lax.stop_gradient(x) + jnp.arange(N)[:, None] * jnp.arange(N)[None, :].)

[pipme]

Thank you for the quick response! It's just a toy example. I will try to wrap the code and share a fully runnable example.

[pipme]

I would guess that the reason execution is slow in the above code without jit, even without any autodiff, is that the updates to J_no_grad are not in-place. They only become in-place updates when under a jit

Yes it takes 2s even without any autodiff. I am writing like this as for numpy it's ok and very fast for the real case. But for Jax we don't have in-place updates. But I just realised maybe I can use a nested list for J_no_grad instead to avoid the at[i,j].set updates?

[pipme]

Hi Matthew, thanks for the help! After further inverstigations it turns out I could vectorize my computations and use lax.stop_gradient for several variables. One curious question, why Jax doesn't have something equivalent to with torch.no_grad()? Is it just not implemented or are there some technical issues? IMO this with statement could be very useful. @mattjj