Repository navigation
Regression test and warning for the Diffrax fixed-step checkpoint-adjoint defect - #68
Open
rogeriojorge wants to merge 1 commit into
Open
rogeriojorge wants to merge 1 commit into
rogeriojorge wants to merge 1 commit into
Conversation
With adaptive_time_step=False and many saves on step ends, Diffrax 0.7.2 RecursiveCheckpointAdjoint gives gradients a few percent off (2.5% for 201 saves at dt = 0.01). Add a strict-xfail regression test against ForwardMode, a passing test of the exact workaround (checkpoints = max_steps) and a docstring warning.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #44 (base
agent/autodiff-plan); no CI runs on a stacked base, tests were run locally.Problem. With
adaptive_time_step=False, Diffrax 0.7.2RecursiveCheckpointAdjointreturns wrong gradients when manySaveAt(ts)times fall exactly on step ends, whichlinspace(0, t_max, timesteps)usually does. Forward values are right. On the Orszag-Tang test case (201 saves, dt = 0.01, max_steps = 220) the default adjoint is 2.5% off ForwardMode. The existing tests use at most 5 saves and miss it. Upstream report: patrick-kidger/diffrax#788Change.
test_fixed_step_many_saves_default_adjoint_matches_forward_mode: strict xfail, so it fails loudly once Diffrax is fixed.test_fixed_step_many_saves_one_checkpoint_per_step_is_exact: the workaroundRecursiveCheckpointAdjoint(checkpoints=max_steps)matches ForwardMode to 1e-9.adjointwith the workarounds.Measured cost (same case, laptop CPU):
No automatic default.
dtandt_maxare traced inside the jittedsimulation, so the step count is not known at trace time.checkpoints=max_stepswith the defaultmax_steps=1e6would store a million states, and DirectAdjoint memory also grows withmax_steps. So the PR documents the workaround and does not change the default.Tests:
pytest -n 2 testsgives 19 passed, 1 xfailed.