[mshettyFE]

Context

I have a physics code based on diffrax which solves its' central system of dif. eqs. via Kvaerno5. For performance reasons, I decided to break up the integration region into two diffrax solves, where at early times I used a Kvaerno5 to handle the stiffness of the equations, and at some cutoff, I switch over to a lighter solver like Tsit5 when the stiffness is gone. This make the code run 2x faster without any noticeable accuracy loss. Hurray!

The problem arose when I tried scaling up the code to run in batches on a H100. Previous, I was able to get a certain number of instances running, but after the performance change, the H100 OOM'd on the same number of instances. I had Claude dump the XLA memory allocation via .lower().compile().memory_analysis() and matched each allocation to a particular source code line number via the metadata, and the biggest allocations were happening in the loop function in integrate module during the t0==t1 branch. All of the allocations in my user-space code were sub-dominant to these allocations.

Cause

The previous behavior used lax.cond to toggle what save state behavior to implement for a given slice. From my understanding, lax.cond allocates the copy buffer for both branches unless it can can guarantee that one of the branches is dead. For my use case, t0 is not consistent across all the solvers since the transition point to the 2nd solver changes depends on the simulation parameters. Hence, the dead code cleanup pass does not happen, and you get stuck with higher temporary memory usage from a branch which never executes.

Solution

Replace the previous _save_if_t0_equals_t1 with a version that uses jnp.where and in-place update to avoid the copy buffer.

Benchmarking

As a stripped down example, I wrote benchmarks/integrate_memory_usage.py that can be used to A/B test the change. It also passes all of diffrax's existing unit tests, the most important of which are _integrate and save_at_solution. I also wrote a regression test prevent the memory usage from slipping again.

I tried my best to keep things generic, but I might be missing some edge case which makes this not scale to all diffrax users. Please LMK what you think. I would also be happy to provide more concrete numbers if desired.