[nstarman]

Dense solutions allocate interpolation buffers up to the solve's max_steps, which can be much larger than needed for later evaluation. Add diffraxtra.subsample(sol, max_steps=...) to compact that storage after integration.

If the accepted steps fit, preserve the original knots and interpolation exactly while trimming padding. Otherwise, sample the original dense curve and its derivatives on a uniform grid and fit cubic Hermite pieces. The returned object remains a diffrax.Solution, with only its interpolation replaced using eqx.tree_at; saved values, statistics, and solver state retain their original meaning. A shared packed coefficient layout lets lax.cond select trimming or resampling under JIT, including mixed batches.

Resampling introduces additional interpolation error and does not inherit the solver's tolerance guarantee. The documentation explains the convergence check, static size requirement, and scope for smooth ODEs.

Validation:

• Regression tests and public API doctests cover exact trimming, convergence, buffer sizes, forward/backward solves, PyTrees, gradients, JIT, batching, early termination, and trivial intervals.
• The full 194-test suite passes in both the locked environment and a minimum-version environment using Diffrax 0.6.0, Equinox 0.11.5, and JAX 0.4.35.
• Ruff, mypy, pylint, and pre-commit checks pass. Analytic force checks pass, and the harmonic dense solution agrees with its analytic orbit within 1e-7 kpc.
• A self-contained example compares harmonic and eccentric Kepler orbits using handwritten JAX potential functions and acceleration computed with -jax.grad(potential). The example only adds matplotlib; it has no galax or unxt dependency. Solving, checking the compact solutions, and plotting are separate helpers, with a small main that assembles the figure. Harmonic trimming from 4096 to 32 slots preserves evaluations exactly while reducing dense storage from 3.18 MB to 24.9 KB. The Kepler example shows convergence with increasing budgets and exact preservation once all 62 accepted steps fit.

Orbit comparisons:

Left: original and compacted orbits. Right: position error relative to the original dense solution. Increasing the resampling budget reduces interpolation error; once all accepted steps fit, trimming preserves the original curve exactly. Exact-zero errors are plotted at a floor of 1e-15 kpc on the logarithmic axes.

[codecov]

Codecov Report

❌ Patch coverage is 95.23810% with 4 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@86d6566). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/diffraxtra/_src/subsample.py	95.18%	4 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main      #57   +/-   ##
=======================================
  Coverage        ?   97.31%           
=======================================
  Files           ?        6           
  Lines           ?      261           
  Branches        ?        0           
=======================================
  Hits            ?      254           
  Misses          ?        7           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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