Skip to content

fix(adjoint): keep saves that land on a step end - #789

Open
Cherith1222 wants to merge 2 commits into
patrick-kidger:mainfrom
Cherith1222:fix/788-checkpoint-step-end-saves
Open

Cherith1222 wants to merge 2 commits into
patrick-kidger:mainfrom
Cherith1222:fix/788-checkpoint-step-end-saves

Conversation

@Cherith1222

Copy link
Copy Markdown

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

Fixes #788

Drafted with Cursor.

A recomputed step end can sit two ULPs below the forward value, so a sample that lands on that end is dropped and its cotangent is zero.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant