Break a fori_loop

[aespielberg]

Is there any way to break out of a fori_loop early? I understand why jax would need to precompile a maximum number of loop iterations, but is there any way to have it run for only a subset of those loop iterations at runtime?

[jakevdp]

There's not any way to break a fori_loop; if you need to break based on some condition you can do so using a while_loop. It might look something like this:

from jax import lax

def func(i, x):
  return x + i
print(lax.fori_loop(lower=0, upper=10, body_fun=func, init_val=0))
# 45

def body_fun(carry):
  i, x = carry
  return i + 1, func(i, x)

def cond_fun(carry):
  i, x = carry
  break_condition = (i == 20)
  return ~break_condition & (i < 10)

print(lax.while_loop(cond_fun, body_fun, init_val=(0, 0))[1])
# 45

[aespielberg]

FWIW, I am pretty sure that jnp.cond is actually evaluating both branches in a different context (which is now becoming a bottleneck for this research project as it is using a tooooon of memory): #8409

[shailesh1729]

Kindly allow me to digress a little bit. Differentiability is not a key concern in many numerical linear algebra routines which aim to solve problems like A x = b for underdetermined systems, computing a few singular values for truncated SVD based algorithms, etc. While I can code such algorithms using JAX but static size restrictions, lack of early exit in fori_loop, requiring that the shape of algorithm state in fori_loop or while_loop iterations cannot change, means that their natural implementation cannot be put under jit. And without jit, JAX is usually quite slow. One has to go through hackish ways to make them work under jit. But it doesn't scale well.

Consider a singular value thresholding algorithm for matrix completion problem. I may have a huge matrix X which is quite sparse (I can now store X efficiently using jax.experimental.sparse). One way is to compute dense SVD (supported by JAX) and then do thresholding. But this is wasteful if X is of the order of 100s of thousands but the actual singular values of concerns are only the top few hundred. My ideal would be to compute a small (say r) number of singular values each time, which may require a variable number of bidiagonalization steps to converge and stop when enough singular values have stabilized up to the specified threshold. Such an algorithm is next to impossible to code and compile under the current jit. I would have to do a worst-case estimate for the number of iterations required for the algorithm to converge so that I can allocate space for storing the intermediate Ritz vectors and other variables as part of the algorithm state. I had recently coded a jittable truncated SVD routine using JAX but I found that it is actually running slower on a GPU than the reference implementation in MATLAB running on a CPU. It is still a significant improvement from full SVD though for huge+sparse matrices.

I do understand that jit, grad, pmap, vmap, jvp, primitives are designed to work interchangeably which imposes these restrictions. But a separate version of jit which doesn't require differentiability on the input function, but only focuses on generating efficient code for the underlying hardware, may open up JAX to many other areas. That special jit could potentially allow dynamic output shapes (depending on function arguments or early stopping) as well as dynamic slices in the intermediate steps.

Please forgive me if some of these comments seem immature as I am still quite new to JAX.

[aespielberg]

Well, I should note that i do care about differentiability, and I think adding a second jit is probably a huge task. I think that there are a ton of hidden catches in what you are proposing and I don't think most of this is related to the thread at hand. If you want to propose jit2, I think you should open a feature request.

What I am asking about is something that should be completely feasible in terms of differentiability and the current jit framework (thought maybe with a bit of work), assuming both jax.lax.cond already works, and given the fact that a differentiable for loop is essentially composing the same function that only has to be composted into itself repeatedly with an input jvp for backwards diff.

[patrick-kidger]

For those following this topic / those who come across this topic.

Diffrax (a pure-JAX library of ODE/SDE/etc. solvers) implements such an "early-exit-but-differentiable-bounded-iteration-loop" as bounded_while_loop, here:
https://github.com/patrick-kidger/diffrax/blob/2b4e4d863c15abc7143919bac7825090bbfe50be/diffrax/misc/bounded_while_loop.py

For those curious about the technical details and caveats, have a look at my reply here:
#2139 (comment)

[aespielberg]

Awesome, thank you so much! I will try this out soon!

[patrick-kidger]

Sorry, meant to post as a proper answer, not just in the comments:

#8375 (reply in thread)

[TL;DR links and explanation for an implementation of a reverse-mode differentiable bounded while loop.]

[njwfish]

Is there any possibility of building this into base jax? It is extremely useful!