From e7e42a20129efeaced96b1a64357d5c68afc9246 Mon Sep 17 00:00:00 2001 From: Mihir Shetty Date: Tue, 20 Oct 2026 13:19:18 -0400 Subject: [PATCH 1/2] Added toy memory bench mark to demo temporary memory usage --- benchmarks/integrate_memory_usage.py | 123 +++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 benchmarks/integrate_memory_usage.py diff --git a/benchmarks/integrate_memory_usage.py b/benchmarks/integrate_memory_usage.py new file mode 100644 index 00000000..f76bf3d4 --- /dev/null +++ b/benchmarks/integrate_memory_usage.py @@ -0,0 +1,123 @@ +"""does the t0 == t1 identity cond in diffrax's loop() cost a +copy of the save buffer? + +Toy Workload: B trajectories of an n-dim linear ODE y' = A y, Tsit5 + PID. +Reports XLA temp memory against the size of one batched save buffer +(B * n_ts * n * 8 bytes for --saveat ts, B * max_steps * n * 8 bytes for --saveat steps) + +How t0 reaches the cond (--t0): + const vmap over y0, t0 a Python constant: XLA can fold t0 == t1 and drop the cond + shared vmap over y0, t0 a traced scalar shared by all trajectories: the inner cond + predicate is unbatched, so it stays a real cond + batched vmap over (y0, t0): the inner cond predicate is batched, so vmap turns it + into a select over the whole save state + novmap no vmap, the batch is solved as one (B, n) ODE with a traced scalar t0 + +Compare temp memory between cases with the same --saveat and --consume flags, not +against an absolute ratio: other loop state also lands in temp. + + python repro_t0t1.py [--t0 {const,shared,batched,novmap}] [--saveat {ts,steps}] + [--consume] [--B B] [--n N] [--n-ts N_TS] + +""" + +import argparse +import json + +import jax +import jax.numpy as jnp + + +jax.config.update("jax_enable_x64", True) + +import diffrax + + +parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter +) +parser.add_argument( + "--consume", action="store_true", help="use sol.ys downstream of the solve" +) +parser.add_argument( + "--t0", + choices=["const", "shared", "batched", "novmap"], + default="const", + help="how t0 is passed in (see above)", +) +parser.add_argument( + "--saveat", choices=["ts", "steps"], default="ts", help="SaveAt(ts=...) or steps" +) +parser.add_argument("--B", type=int, default=1000, help="number of trajectories") +parser.add_argument("--n", type=int, default=100, help="state dimension") +parser.add_argument("--n-ts", type=int, default=1000, help="number of save points") +cli = parser.parse_args() + +consume = cli.consume +B, n, n_ts = cli.B, cli.n, cli.n_ts +max_steps = 4096 + +A = -jnp.eye(n) + 0.1 * jax.random.normal(jax.random.PRNGKey(0), (n, n)) / jnp.sqrt(n) +ts = jnp.linspace(0.0, 10.0, n_ts) +if cli.saveat == "ts": + saveat = diffrax.SaveAt(ts=ts) +else: + saveat = diffrax.SaveAt(steps=True) + + +def solve(y0, t0): + sol = diffrax.diffeqsolve( + diffrax.ODETerm(lambda t, y, args: y @ A.T), # y is (n,) or, for novmap, (B, n) + diffrax.Tsit5(), + t0=t0, + t1=10.0, + dt0=0.01, + y0=y0, + saveat=saveat, + stepsize_controller=diffrax.PIDController(rtol=1e-8, atol=1e-8), + max_steps=max_steps, + ) + ys = sol.ys + assert ys is not None + if consume: # use the saved ys downstream to prevent XLA optimizing it away + first_half = jnp.arange(len(ys)) < len(ys) // 2 + mask = first_half.reshape((-1,) + (1,) * (ys.ndim - 1)) + ys = jnp.where(mask, ys, 2.0 * ys) + return ys + + +y0 = jax.random.normal(jax.random.PRNGKey(1), (B, n)) +t0 = jnp.array(0.0) +if cli.t0 == "const": + f, args = jax.jit(jax.vmap(lambda y: solve(y, 0.0))), (y0,) +elif cli.t0 == "shared": + f, args = jax.jit(jax.vmap(solve, in_axes=(0, None))), (y0, t0) +elif cli.t0 == "batched": + f, args = jax.jit(jax.vmap(solve)), (y0, jnp.zeros(B)) +else: + f, args = jax.jit(solve), (y0, t0) +m = f.lower(*args).compile().memory_analysis() +assert m is not None +buf = B * (n_ts if cli.saveat == "ts" else max_steps) * n * 8 +ys = f(*args) # also run it. Answer shouldnt change +print( + json.dumps( + dict( + consume=consume, + t0=cli.t0, + saveat=cli.saveat, + B=B, + n=n, + n_ts=n_ts, + device=str(jax.devices()[0]), + save_buffer_GB=round(buf / 1e9, 3), + temp_GB=round(m.temp_size_in_bytes / 1e9, 3), + out_GB=round(m.output_size_in_bytes / 1e9, 3), + temp_in_save_buffers=round(m.temp_size_in_bytes / buf, 2), + # dirty check to make sure output doesnt degrade; steps pads unused rows + # with inf, so only sum the finite entries + checksum=float(jnp.sum(jnp.where(jnp.isfinite(ys), ys, 0.0))), + ) + ), + flush=True, +) From 137f6bcee2b9ccea10d31c77e6f639be465ea8bc Mon Sep 17 00:00:00 2001 From: Mihir Shetty Date: Thu, 22 Oct 2026 17:51:23 -0400 Subject: [PATCH 2/2] Changed t0==t1 behavior of _integrate.loop to avoid unnecessary copy buffer allocation --- diffrax/_integrate.py | 44 ++++++++++++++++++++---------------- test/test_saveat_solution.py | 42 ++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 19 deletions(-) diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py index 5441e0a9..aa16b81f 100644 --- a/diffrax/_integrate.py +++ b/diffrax/_integrate.py @@ -820,28 +820,34 @@ def unsave(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState: unsave, saveat.subs, save_state, is_leaf=_is_subsaveat ) - def _save_if_t0_equals_t1(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState: - if subsaveat.ts is not None: - save_state = _save( - t0, yfinal, args, subsaveat.fn, save_state, repeat=len(subsaveat.ts) - ) - return save_state - # if t0 == t1 then we don't enter the integration loop. In this case we have to # manually update the saved ts and ys if we want to save at "intermediate" # times specified by saveat.subs.ts - save_state = jax.lax.cond( - eqxi.unvmap_any(t0 == t1), - lambda __save_state: jax.lax.cond( - t0 == t1, - lambda _save_state: jtu.tree_map( - _save_if_t0_equals_t1, saveat.subs, _save_state, is_leaf=_is_subsaveat - ), - lambda _save_state: _save_state, - __save_state, - ), - lambda __save_state: __save_state, - save_state, + # + # For subsaveat states, only 't0' row has been written. Hence, the 'ts' rows start + # at a static offset. Instead of lax.cond, we can statically calculate what slice + # of the state to update. This coaxes XLA optimizer to not emit an extra copy buffer + # of the full difeq state needlessly + # + def _save_if_t0_equals_t1(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState: + if subsaveat.ts is None: + return save_state + pred = t0 == t1 + start = 1 if subsaveat.t0 else 0 + stop = start + len(subsaveat.ts) + + def _update(buffer, value): + return buffer.at[start:stop].set(jnp.where(pred, value, buffer[start:stop])) + + ts = _update(save_state.ts, t0) + ys = jtu.tree_map(_update, save_state.ys, subsaveat.fn(t0, yfinal, args)) + save_index = jnp.where(pred, stop, save_state.save_index) + return eqx.tree_at( + lambda s: [s.ts, s.ys, s.save_index], save_state, [ts, ys, save_index] + ) + + save_state = jtu.tree_map( + _save_if_t0_equals_t1, saveat.subs, save_state, is_leaf=_is_subsaveat ) def _save_t1(subsaveat, save_state): diff --git a/test/test_saveat_solution.py b/test/test_saveat_solution.py index d1aa6037..db9ac6bf 100644 --- a/test/test_saveat_solution.py +++ b/test/test_saveat_solution.py @@ -428,6 +428,48 @@ def _solve(tf): assert tree_allclose(yc, true_yc) +@pytest.mark.parametrize("batch_t0", [False, True]) +def test_t0_eq_t1_no_save_buffer_copy(batch_t0): + """Handling t0 == t1 must not cost an extra copy of the save buffer. + Compares XLA's planned temporary memory against a constant t0, for which the + t0 == t1 check is constant-folded away. `throw=False` keeps the `error_if` on the + solution, which can cost a copy of its own on GPU, out of the comparison. + """ + num_batch, num_ts, size = 16, 200, 16 + ts = jnp.linspace(0.0, 1.0, num_ts) + + def _solve(y0, t0): + sol = diffrax.diffeqsolve( + diffrax.ODETerm(lambda t, y, args: -y), + diffrax.Tsit5(), + t0=t0, + t1=1.0, + dt0=0.1, + y0=y0, + saveat=diffrax.SaveAt(ts=ts), + # On CPU, the old cond-based handling only cost a copy with an adaptive + # controller; with a constant step size XLA happened to fold it. + stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-6), + throw=False, + ) + return sol.ys + + def _temp_bytes(fn, *args): + analysis = jax.jit(fn).lower(*args).compile().memory_analysis() + if analysis is None: + pytest.skip("memory analysis is unavailable on this backend") + return analysis.temp_size_in_bytes + + y0 = jnp.ones((num_batch, size)) + constant = _temp_bytes(jax.vmap(lambda y: _solve(y, 0.0)), y0) + if batch_t0: + traced = _temp_bytes(jax.vmap(_solve), y0, jnp.zeros(num_batch)) + else: + traced = _temp_bytes(jax.vmap(_solve, in_axes=(0, None)), y0, jnp.array(0.0)) + save_buffer = num_batch * num_ts * size * y0.dtype.itemsize + assert traced - constant < save_buffer / 2 + + def test_trivial_dense(): term = diffrax.ODETerm(lambda t, y, args: -0.5 * y) y0 = jnp.array([2.1])