Future of for_loop?

[lockwo]

The for_loop (distinct in what it allows inside from the fori_loop) has been in the code base for a while (https://github.com/jax-ml/jax/blob/main/jax/_src/lax/control_flow/for_loop.py#L82). Are there any plans to develop or delete it? It seems pretty hidden currently (only 3 GitHub issues, and it shows up in the docs in one place https://jax.readthedocs.io/en/latest/pallas/tpu/details.html), but I'm curious as to the plans there because I think it is interesting.

[jakevdp]

@sharadmv may be able to say more about the plans (tracked in issue #10982).

[mattjj]

Thanks for asking!

As you may know, we added array references and effectful operations on them to JAX for the purpose of decomposing scan into something more flexible. Then that machinery turned into the basis of Pallas, and Pallas has basically sidetracked us ever since. We kept meaning to get back to landing for_loop...

Now the plan is to unbundle things further: where for_loop builds in the introduction of array references (in the sense that they are introduced as lambda-binders in the body function), we plan to separate that out into top-level mutable array creation. With that, the loop construct itself gets even simpler. (See this test file to get a flavor.)

That further unbundling requires some upgrades to the autodiff system, which we're currently doing under the moniker "direct linearize" (as in JAX_USE_DIRECT_LINEARIZE=1 etc). (The original for_loop bundled those autodiff changes under its autodiff rule, effectively.) Once we land that, plus perhaps a few other changes, we'll surface the improved loop construct.

[lockwo]

Awesome, thanks for the info!