Test that control checkpoints survive recompute (fixes #5082) - #5093
Test that control checkpoints survive recompute (fixes #5082)#5093sghelichkhani wants to merge 3 commits into
Conversation
The clear-down assertions in _check_forward, _check_recompute, and _check_reverse swept every block variable and required _checkpoint to be None after recompute. That codified the pre-fix behaviour where the checkpoint manager would silently clobber the control's checkpoint during forward replay. Now that the manager goes through the BlockVariable setter and respects the is_control guard, the control's block variable legitimately retains the user-supplied value across replays. The helpers now take an optional list of controls and skip those block variables. Also add a regression test for #5082: a 4-step model with J = sum(m**2) evaluated at m0 = 2 must have derivative 16 (the bug produced 8 because the adjoint replay saw the original m = 1).
|
Thanks for tackling this! We need to be running this branch against the corresponding pyadjoint branch though, so ou'll need to change this line in the Line 39 in 6386785 to: "pyadjoint-ad @ git+https://github.com/dolfin-adjoint/pyadjoint.git@sghelichkhani:sghelichkhani/fix-control-checkpoint-clearing",
FYI I've not used an alternative branch from a fork before so you might have to change it to point to your fork directly instead of the main repo. Arguably this should go into |
|
I've set up a draft PR #5157 which has @sghelichkhani 's commit cherry-picked on top of release and testing against the sghelichkhani/fix-control-checkpoint-clearing branch of pyadjoint. So what would be the order of things (assuming tests are passing): we get dolfin-adjoint/pyadjoint#257 merged into pyadjoint master, and create a patch version?, then merge this into Firedrake release? |
Great, thanks.
I think the order would be (@connorjward?):
I can put it on the agenda for the meeting this afternoon. |
Yeah this seems like the right thing. Don't make a patch release if you can avoid it (creates churn). |
The _control_bvs helper accepted either a Control or the underlying overloaded variable, which is more flexible than the call sites need. Standardise on the overloaded variable, which is available at every call site (including before the Control is constructed), and update the two multistep call sites that were passing a Control.
Point the dependency at sghelichkhani/fix-control-checkpoint-clearing so CI exercises these test changes against the matching pyadjoint fix. Reverts to the released pyadjoint once dolfin-adjoint/pyadjoint#257 is merged and tagged.
|
Thanks for addressing my comment @sghelichkhani. Because this is a bugfix we want to merge it into |
|
Thanks for getting 257 in and sorting the order out. Happy to go with whatever you decide on timing and how it reaches main, so please don't cut a point release just for me. Nothing in g-adopt is actually failing on this, so I'd rather it went out with the next normal release than add churn. The only related thing still open on my end is pyadjoint #248 (the SingleMemoryStorageSchedule clearing fix), and that case is xfailed for now, so nothing's going red while we wait. Just flagging it's on my radar, but I'm fine for it to sit. |
I'll admit that I am very unfamiliar with the checkpointing in general, but doesn't that do the same thing you are trying to fix with the current PRs? Accessing the If you put it on the meeting agenda for next week then we can take a look at it. |
Companion to dolfin-adjoint/pyadjoint#257, which fixes the underlying bug in
CheckpointManagerwhere direct writes toBlockVariable._checkpointbypassed theis_controlguard and wiped the user-supplied control value during forward replay. With that pyadjoint fix in place the control's checkpoint now legitimately survives across replays, which makes two changes necessary here.The first is that
_check_forward,_check_recomputeand_check_reverseintests/firedrake/adjoint/test_burgers_newton.pywere asserting that every block-variable checkpoint must beNoneafter replay. That assertion codified the buggy behaviour — under the fix, the control's block variable correctly retains its checkpoint between calls torf(new_value)andrf.derivative(). The helpers now take an optional list of controls (accepting eitherControlobjects or the underlying overloaded variables) and skip those block variables in the clear-down sweeps. All existing call sites intest_burgers_newton.pyandtest_checkpointing_multistep.pyare updated to thread the controls through.The second is a regression test
test_control_value_survives_recomputeintest_checkpointing_multistep.pythat captures Steph's MFE from #5082 directly: four timesteps ofJ = sum_k m**2, evaluated atm0 = 2, must give a derivative of16. Before the pyadjoint fix this returned8because the adjoint replay saw the stale underlyingm = 1.Depends on dolfin-adjoint/pyadjoint#257 — without it the new regression test will fail.