perf(vlasov1d): halve reverse-mode residuals in the spectral velocity push

[joglekara]

What

The spectral velocity push is

irfft(exp(-1j * kv * dt * accel) * rfft(f))

accel comes from the field solve, so it is traced, and reverse-mode has to keep two (nx, nv//2+1) complex128 arrays alive per call: the spectrum (needed by the phase gradient) and the phase factor (needed by the spectrum gradient). That is two residual arrays for every one distribution function advected, and SixthOrderHamIntegrator issues six pushes per timestep.

This routes the phase multiply through a custom_vjp that saves only the spectrum and the small generators (kv, dt, accel), and rebuilds the phase factor on the backward pass.

Measured

Residual bytes read straight off the linearized jaxpr, so the numbers are a property of JAX's partial evaluation rather than of any one backend. At nx=64, nv=4096 (the bump-on-tail.yaml velocity resolution):

	residual / push
before	4.229 MB
after	2.115 MB

Exactly the phase factor dropped — 50.0%. Under the sixth-order integrator that is ~25 MB/step of edfdv residual against a 2.1 MB state, so the per-step figure goes to ~12.7 MB.

Correctness

The primal and both gradients come out bit-identical to the expression they replace. The tests assert equality, not a tolerance — a silent gradient change in an optimization codebase is the failure mode worth guarding against, and bit-identity means no downstream behaviour can shift.

Two choices make that guarantee cheap to hold:

• The wrapper covers only the phase multiply, not the surrounding transforms. rfft/irfft are linear, so autodiff already stores nothing for them; leaving them untouched keeps the FFT conventions (including the Nyquist bin, which a spectral shift does not treat as a pure rotation) exactly as they were.
• The backward rule differentiates a recomputation of the same primal with jax.vjp rather than hand-deriving the complex-multiply adjoint. Exact by construction rather than by argument.

Scope

SpaceExponential and HouLiFilter are deliberately left alone. Their factors are built from constants — the v-grid times dt, and the fixed filter kernel — so XLA hoists them out of the step and there is no per-step residual to save. Only VelocityExponential has a traced phase. (If dt ever becomes traced, e.g. adaptive stepping, SpaceExponential would become a candidate too.)

Tests

New tests/test_vlasov1d/test_pusher_spectral_vjp.py — 6 tests:

• forward is bit-identical to the inline expression
• gradients are bit-identical through both the f and accel paths
• the phase factor is provably no longer a residual (asserted in bytes, not inferred)
• end-to-end through the real VelocityExponential class
• composes with shard_map — parallel: [x] (used by the large IAW configs) wraps push, so the custom VJP has to survive that. Runs in a subprocess since faking devices needs XLA flags set before JAX initializes; there was no prior sharding coverage here.

Full suite: 113 passed, 7 deselected, 1 xpassed. The xpass is the pre-existing self-marked flaky test_landau_damping ("Fails mysteriously with float64"), unrelated.

CI needs no change — the existing paths filter already covers adept/_vlasov1d/** and tests/test_vlasov1d/**.

Tradeoff worth knowing

This buys memory with compute: one extra elementwise exp on the reverse pass. On CPU that measures ~7% slower backward for the 50% memory. On GPU these kernels are bandwidth-bound, so recomputing an exp should beat the 2 MB round trip it replaces — but that wants confirming on real hardware before anyone leans on it. A GPU A/B benchmark script (peak memory + backward time, monkeypatching the helper back to the plain expression) is available separately; it was kept out of the tree since there is no benchmarks/ dir. If it comes back negative for production grid sizes, the change is one line to revert at the call site.

The payoff is not the 2 MB by itself — it is that halving per-step residuals lets RecursiveCheckpointAdjoint afford more checkpoints at the same memory ceiling, which is the dominant backward cost on the long IAW runs.

---

Generated by Claude Code

[ergodic-ludwig]

Review summary

Verdict: Changes requested

The reverse-mode memory optimization and its reverse-mode/shard_map coverage are well motivated, and the relevant CI checks pass. However, introducing custom_vjp in the unconditional velocity-push path removes forward-mode autodiff support. That is a blocking compatibility regression for this solver because the repository documentation explicitly identifies jax.jvp as the supported differentiation mode for the parallel: ["v"] configuration. See the inline comment for the requested fix and regression test.

[ergodic-ludwig]

Blocking: jax.custom_vjp makes _shift_spectrum unavailable to forward-mode autodiff (jax.jvp/jacfwd), so routing every VelocityExponential.push through it regresses a previously supported differentiation path. This matters here because docs/source/solvers/vlasov1d/overview.md explicitly directs gradient users to forward mode for the parallel: ["v"] configuration where reverse mode is unsupported. Please preserve a JVP-capable path (for example, use a custom primitive/custom derivative strategy that supports both modes, or conditionally retain the original expression for forward mode) and add an end-to-end jax.jvp regression test through VelocityExponential.