[Cherith1222]

Why

A recursive checkpoint adjoint drops the derivative of a sample that lands on a step end. Recomputing that step rounds the end two ULPs early, so the save is skipped and the sample adds nothing to the gradient.

What changed

• The save loop keeps a time that sits within two ULPs after the step end on the current step.
• A regression compares this gradient with the direct adjoint when every save lies on a constant step.

Scope

• Covers output times that fall on a step boundary in a checkpointed solve.
• Leaves adaptive step-size control, and saves that sit clearly inside a step, alone.

Blast Radius

This touches solves that record output on step ends and then differentiate them. Forward samples were already correct. The wider test only pulls in a time that rounding had pushed just past the recomputed end.

Verification

• The new test failed before the change, 6.924 against 7.174, and passes after it.
• The oscillator in RecursiveCheckpointAdjoint returns wrong gradients (~2%) when SaveAt(ts) times coincide with ConstantStepSize step ends; result depends on checkpoint count #788 now matches the direct adjoint. Both gradients are 331.67329043219127. The closed form is 331.6732904321968.
• test/test_adjoint.py without the slow test: 11 passed. test/test_saveat_solution.py: 26 passed.

Fixes #788

Drafted with Cursor.

[salma2vec]

I checked 0523b8c against #791 too. The reported float32 rotation ends with inf on main; this head returns [0.7316889763, 0.6816388965], within 1.36e-7 of [cos(0.75), sin(0.75)], with the same 108 accepted and 45 rejected steps.

I also checked real/complex states in both precisions and directions, early event/step-limit padding, and the constant-field derivative under recursive/direct reverse mode and forward mode. The 17 focused cases pass with numerical tolerances. A small regression for the terminal-jump case could also cover #791 here:

def test_save_at_terminal_jump():
    ts = jnp.array([0.25, 0.5, 1.0], dtype=jnp.float32)
    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(lambda t, y, args: 2 * jnp.ones_like(y)),
        diffrax.Tsit5(), ts[0], ts[-1], ts[-1] - ts[0],
        jnp.array(1.0, dtype=jnp.float32),
        saveat=diffrax.SaveAt(ts=ts),
        stepsize_controller=diffrax.ClipStepSizeController(
            diffrax.PIDController(rtol=1e-5, atol=1e-5),
            jump_ts=ts[-1:],
        ),
    )
    assert sol.result == diffrax.RESULTS.successful
    assert jnp.array_equal(sol.ts, ts)
    assert jnp.allclose(sol.ys, 1 + 2 * (ts - ts[0]), rtol=1e-6, atol=1e-6)

Checked on CPU with JAX 0.10.2, Equinox 0.13.8, and Python 3.11.15.