Add checkpoints parameter for reverse-mode adjoint control

[FFroehlich]

Summary

This PR adds a checkpoints parameter to the JAX PETAb simulation interface, allowing users to control the number of checkpoints used by diffrax's RecursiveCheckpointAdjoint when computing gradients. This provides fine-grained control over the memory-computation tradeoff during backpropagation.

Key Changes

• Added checkpoints: int | None = None parameter to run_simulation() method
• Added checkpoints: int | None = None parameter to run_simulations() method (both the class method and module-level function)
• Updated RecursiveCheckpointAdjoint() instantiation to pass the checkpoints parameter
• Added checkpoints to the eqx.filter_vmap configuration to exclude it from vectorization
• Added comprehensive docstring documentation explaining the parameter's behavior and tradeoffs

Implementation Details

• The checkpoints parameter is passed through the call chain from the public API down to the diffrax adjoint configuration
• When None (default), diffrax automatically selects ~sqrt(2 * max_steps) checkpoints, preserving backward compatibility
• The parameter only affects reverse-mode adjoints (used for llh and chi2 return values); other return values use DirectAdjoint which ignores this setting
• Larger checkpoint values reduce backward-pass recomputation at the cost of increased memory usage and compilation time, particularly beneficial for models with long/stiff dynamics that exceed the default checkpoint count

https://claude.ai/code/session_01EbUmB8BLG5E36KKoS22ttG

[FFroehlich]

Benchmark report behind this PR

This PR exposes an opt-in checkpoints argument (default None = unchanged behavior). Below is the investigation that motivated it — the question was whether explicitly setting the adjoint checkpoints (to max_steps, max_steps/2, max_steps/4) gives a performance advantage. Short answer: only for long/stiff trajectories, so the default is deliberately left unchanged and the tuning is opt-in.

Methodology note

Reverse-mode value_and_grad(run_simulations), atol=rtol=1e-8, CPU, float64. Models imported fresh from benchmark_models_petab through the JAX backend. checkpoints values were injected to compare against the current default; the committed change wires the same value through run_simulations → JAXProblem.run_simulations → JAXProblem.run_simulation into RecursiveCheckpointAdjoint(checkpoints=...). Gradients/llh matched the default within rtol=1e-5 for every setting.

---

What checkpoints controls

The gradient path uses diffrax.RecursiveCheckpointAdjoint() with checkpoints=None. With None, equinox's checkpointed while_loop picks a default of ≈√(2·max_steps) checkpoints (binomial/online-treeverse). More checkpoints ⇒ fewer forward-segment recomputations in the backward pass ⇒ faster, at the cost of memory + compile time.

Key consequence: raising checkpoints only helps when the actual number of solver steps exceeds the checkpoint budget. If the trajectory already fits, every step is stored and there is zero recomputation to eliminate. For AMICI defaults, √(2·max_steps) is: max_steps=2**13 → 127; 16384 → 181; 2·10**5 (gradient test) → 632.

PEtab Benchmark Collection

Speedup = min-exec time vs. the default(None) row (higher = faster). ms = max_steps (8192, except Borghans 16384); num_steps = checkpoints sized to the actual step count.

model	nx	npar	steps	default ckpt	num_steps	ms/4	ms/2	ms=max_steps
Bruno_JExpBot2016	7	13	24	127	0.96×	1.02×	1.01×	0.94×
Sneyd_PNAS2002	6	15	68	127	0.90×	0.94×	1.03×	1.02×
Fujita_SciSignal2010	9	19	168	127	1.05×	0.88×	0.64×	0.96×
Fiedler_BMCSystBiol2016	6	22	471	127	1.17×	1.18×	1.19×	1.24×
Boehm_JProteomeRes2014	8	9	664	127	1.25×	1.29×	1.16×	0.81×
Brannmark_JBC2010 (preeq)	9	22	1168	127	1.02×	0.95×	0.82×	0.95×
Elowitz_Nature2000	8	21	2266	127	1.07×	1.14×	1.13×	0.85×
Weber_BMC2015	7	36	3292	127	1.01×	1.18×	1.04×	1.15×
Borghans_BiophysChem1997	3	23	10082	180	1.19×	1.11×	1.24×	1.67×

Reading the table

1. steps ≤ default (Bruno, Sneyd, Fujita): no effect — default already stores the whole trajectory; noise ±5–10%.
2. steps a few× the default (Fiedler, Boehm, Elowitz, Weber): a modest 1.1–1.3× for moderate counts, but pushing to max_steps when max_steps ≫ steps frequently regresses (Boehm 0.81×, Elowitz 0.85×) — over-allocated buffers add overhead.
3. steps ≫ default (Borghans, 10082 vs 180): the one clear win — checkpoints=max_steps gives 1.67×.
4. pre-equilibration-dominated (Brannmark): irrelevant — the steady-state solve uses ImplicitAdjoint, not the checkpointed adjoint.

Cost side

• Compile time jumps from <1–4 s (default; loop built lazily) to 10–60 s per model with an explicit checkpoints (fixed-size loop unrolled). One-time per shape, amortized by JIT caching.
• Memory scales with checkpoints × state_dim.
• checkpoints = max_steps is only feasible when max_steps is modest — the gradient-test max_steps=2·10**5 would allocate 200 000 buffers; Weber's 4·10**7 → 40 M, infeasible.

Cross-check — synthetic Tier-1 models (tests/performance)

max_steps=2**14 (default ckpt 180). Effects look larger here only because these run in µs–ms, where checkpoint bookkeeping is a big fraction of the tiny total:

model	steps	best (setting→speedup)	ms/2	ms/4
LinearDecay	18	num_steps → 1.36×	0.74×	0.84×
ConservationLaw	60	num_steps → 1.99×	0.87×	0.97×
Robertson	197	max_steps → 1.24×	0.96×	0.97×
LotkaVolterra	1780	max_steps → 1.59×	1.19×	1.18×
SingleEvent	26	~neutral	0.12×	0.31×
MultiEvent	58	~neutral	0.15×	0.37×

Event models are catastrophic with large checkpoints (0.12–0.37×): the outer eqxi.while_loop(kind="bounded") wraps a per-segment diffeqsolve, so a giant checkpoint buffer is allocated per segment — a strong argument against ever blanket-setting checkpoints=max_steps.

Bottom line

• checkpoints = max_steps / max_steps/2 / max_steps/4 is not a good general setting. Across the collection it is neutral-to-slightly-harmful for most models, and max_steps specifically regresses several.
• The only robust-win regime is long/stiff trajectories where step count ≫ √(2·max_steps) — Borghans (1.67×), and the synthetic LotkaVolterra (1.59×).
• The principled lever is checkpoints ≈ actual step count (store-all, no recompute), not max_steps — they only coincide when max_steps is snugly sized.
• Recommendation (implemented here): don't change the default; expose checkpoints as an optional argument (default None = current behavior) so long-trajectory models can opt in.

One possible follow-up: the auto-default keys off max_steps, not the actual step count. When max_steps is a large safety cap (e.g. 2·10**5), the default over-provisions for short trajectories and under-provisions for long ones. A step-count-aware heuristic could capture the Borghans win automatically, but that's a larger design change.

---

Generated by Claude Code

[Copilot]

Pull request overview

This PR extends the JAX PEtab simulation API to accept a checkpoints: int | None = None parameter, allowing callers to control how many checkpoints diffrax’s RecursiveCheckpointAdjoint uses during reverse-mode gradient computation (memory vs. recomputation tradeoff).

Changes:

• Added checkpoints parameter to JAXProblem.run_simulation() and JAXProblem.run_simulations().
• Threaded checkpoints through the module-level run_simulations() wrapper down to the diffrax adjoint configuration.
• Updated vmapping configuration to treat checkpoints as non-vectorized/static input.

Comments suppressed due to low confidence (2)

python/sdist/amici/sim/jax/petab.py:1662

• This docstring repeats a specific formula for diffrax’s default checkpoint selection ("~sqrt(2 * max_steps)"). To avoid coupling the docs to an external library’s internal heuristic, consider rephrasing to say diffrax chooses the number of checkpoints automatically when checkpoints=None.

            :meth:`run_simulation` for details. ``None`` (default) preserves the
            previous behaviour (diffrax picks ``~sqrt(2 * max_steps)``).

python/sdist/amici/sim/jax/petab.py:1892

• This public API docstring states a specific default checkpoint formula ("~sqrt(2 * max_steps)"). Since that behavior is owned by diffrax and may vary by version, it would be more robust to describe the default as an automatic choice based on max_steps.

        ``None`` (default) keeps the previous behaviour, where diffrax picks
        ``~sqrt(2 * max_steps)`` checkpoints. Increasing it reduces backward-pass
        recomputation for models whose trajectories take many solver steps

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.83%. Comparing base (8935e72) to head (2be5312).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3217      +/-   ##
==========================================
- Coverage   78.46%   77.83%   -0.63%     
==========================================
  Files         317      317              
  Lines       20974    20974              
  Branches     1483     1482       -1     
==========================================
- Hits        16458    16326     -132     
- Misses       4508     4640     +132     
  Partials        8        8              

Flag	Coverage Δ
cpp	72.07% <ø> (-0.02%)	⬇️
cpp_python	36.67% <ø> (ø)
petab	47.15% <ø> (ø)
petab_sciml	16.21% <ø> (ø)
petab_sciml_benchmarks	14.76% <ø> (ø)
python	70.31% <ø> (ø)
sbmlsuite-jax	?

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
python/sdist/amici/sim/jax/petab.py	91.34% <ø> (ø)

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.