vjp of while_loop with dynamic checkpointing

[dionhaefner]

I was thinking about enabling reverse-mode AD for adaptive solvers (using jax.lax.while_loop internally). As I understand it, the current failure to support that is because JAX cannot know how much memory is needed for the backwards pass before the forward pass is executed, which violates XLA's static memory constraints.

But what if we use a checkpointing scheme that stores a fixed number of checkpoints? In this case we would know in advance how much memory is needed for the backwards pass. I wanted to play around with this but I have no idea how I would go about implementing something like that. I probably need to define a custom vjp for while_loop? Is there existing machinery of jax.remat that I can re-use?

[shoyer]

Take a look at Diffrax's diffrax.misc.bounded_while_loop, which I believe is exactly what you're thinking of. See the implementation notes and #2139 for more relevant discussion and references.

You might also find Diffrax's other machinery for ODE solvers interesting

[dionhaefner]

Thanks, that's an interesting pointer. My proposal is a bit different though, in that it wouldn't need an upper limit on the loop iterations. What I have in mind is this:

1. Specify a fixed number of gradient checkpoints at compile time
2. During forward pass, start filling checkpoint buffer after every iteration
3. If buffer is full, evict old checkpoints according to a policy (this one claims to be optimal)

That way, we could have our cake and eat it, too: open-ended loops, compile-time constant array shapes, O(1) memory use, O(1) forward pass runtime, ~O(log(n)) backward pass runtime. But is it implementable in JAX / XLA? :)