Skip to content

Don't carry the tangent of a clipped step into the next steps - #773

Open
etiferrier wants to merge 2 commits into
patrick-kidger:mainfrom
etiferrier:clip-step-tangent
Open

etiferrier wants to merge 2 commits into
patrick-kidger:mainfrom
etiferrier:clip-step-tangent

Conversation

@etiferrier

@etiferrier etiferrier commented Sep 28, 2026 •

Copy link
Copy Markdown

Fixes #772.

Problem. With traced jump_ts, a step clipped to just before a jump is often tiny, but the tangent of its length, d(jump) - d(t0), is O(1). PIDController proposes the next step as prev_dt * factor, so that tangent is scaled up into every later step time, and it compounds jump after jump. With several close jump times, the derivative with respect to them ends up wrong by many orders of magnitude, in forward and reverse mode.

Fix. In ClipStepSizeController.adapt_step_size, right after a step lands on a jump_ts or step_ts time, the proposed next step keeps its value but takes the tangent of the new t0 (_drop_step_tangent). The following step times then move rigidly with the time just landed on.

Test. test_grad_wrt_close_jump_ts has 7 pulses 0.04 apart, with all 14 edges traced, and runs with and without a step_ts grid. It checks against central finite differences, for RecursiveCheckpointAdjoint, DirectAdjoint and ForwardMode.

Without step_ts With step_ts
main fails: +2.2e14 vs −7.7e-3 fails: +2.5e14 vs −7.7e-3
This PR passes passes

The rest of the test suite behaves as on main.

Downstream check. I ran dynamiqs, which passes pulse edges as jump_ts, on 8 production-like quantum-device workloads:

  • Forward results and step counts are bit-identical, and speed is unchanged.
  • The Jacobian columns with respect to pulse-edge times now match finite differences: e.g. relative error 7.9e45 → 1e-5.

etiferrier and others added 2 commits September 28, 2026 09:20
With traced jump_ts, a step clipped to just before a jump is often very short
while the tangent of its length, d(jump) - d(t0), is O(1). The inner controller
proposes the next step as a multiple of that length, so its tangent scaled the
tangents of all the following step times by (step size / clipped step size),
jump after jump: the derivative with respect to several close jump times was
wrong by many orders of magnitude, in forward and reverse mode (and worse when
step_ts are also present). After a step that lands on a jump_ts or step_ts
time, keep the proposed step's value but give it the tangent of the new t0, so
the following step times move rigidly with the time just landed on.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without the step_ts branch of the fix, the derivative is still off by
1e2-3e10 when step_ts are present; the new parametrization catches it. The
solver tolerance is tightened to 1e-10 so that the step_ts variant separates
clearly (with the fix: <1e-6; without: >1e2).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Gradients with respect to traced jump_ts blow up: the tangent of a clipped step grows through the next steps

1 participant