Fix checkpointed recompute aliasing for state-transfer Function targets

[finsberg]

Summary

FunctionAssignBlock.recompute_component always mutated block_variable.saved_output in place and returned the same object on every recompute. That's required for a Function backing a live dolfinx.fem.DirichletBC (the C++ binding reads that object's array directly, not through the tape), but wrong for an ordinary state-transfer Function (e.g. assign(uh, u_prev) in a time-stepping loop): every timestep's recompute silently overwrote the value an earlier timestep's checkpoint relied on, giving a wrong gradient once a checkpoint schedule actually forces replay (confirmed empirically: Taylor rate 0.70 instead of 2.0 under Revolve, exact 2.0 with no schedule).

This PR went through two designs before landing on the current one — recorded here since both are informative:

1. (shipped first, since replaced) Tag a Function as identity-sensitive via a new Function._ad_bc_backing attribute, set by DirichletBC.__init__, and branch FunctionAssignBlock.recompute_component on it. Worked, but left an implicit, easy-to-forget special case: a boolean flag that silently changes what a generic assign() call does, set from a completely different call site. It also had a documented residual gap — a Function that was both BC-backing and reassigned every timestep still hit the original aliasing bug.
2. Shipped: FunctionAssignBlock isolates unconditionally for every Function target — no tag, no branch. LinearProblemBlock/NonlinearProblemBlock instead track each BC's backing Function (bc.g) as an explicit tape dependency, the same way every other form coefficient already is, and a new sync_bc_values helper (blocks/dirichletbc.py) refreshes bc.g's live array from that pinned dependency's saved_output right before every solve, including during recompute. This closes the residual gap design 1 had to leave open, and needed no new attribute anywhere.

A first attempt at sync_bc_values, keyed off bc.g.block_variable.saved_output directly, was also tried and also empirically wrong — .block_variable always points at bc.g's most recently created BlockVariable, which after the full tape is recorded is simply the last timestep's, regardless of which point in a replay is being recomputed. Reading from the calling block's own pinned dependency (self.get_dependencies()) instead is what's actually position-aware. Recorded in the spec so nobody rediscovers this by bisection.

No change to assign()'s or dirichletbc()'s public signature. The PR #75 RuntimeError guard is untouched. One unrelated cleanup: removed a now-stale xfail(strict=True) on test_snes_time_loop_gradient_is_correct, which started passing due to an earlier, separate fix and was blocking a fully green test_checkpointing.py.

Full design rationale, including the rejected designs and the empirical evidence for each, in .scratch/dirichlet-bc-recompute-identity/spec.md (knowledge repo).

Test plan

• tests/test_checkpointing.py::test_bc_gradient_matches_uncheckpointed (new) — a time-dependent Dirichlet BC under a genuine Revolve schedule, comparing gradients (not just forward values) against the unscheduled run. Fails on the previously-shipped tag-based design with a real ~0.6% gradient mismatch (test_time_dependent_bc_replay never enables an actual schedule, so it couldn't catch this); passes exactly on the current design.
• tests/test_assign.py::test_recompute_does_not_alias_state_across_timesteps — reproduces the original state-transfer aliasing defect directly via assign() chains under a Revolve schedule, no PDE solve
• tests/test_dirichlet_bc.py::test_time_dependent_bc_replay — stays green throughout (the test that would catch a regression toward "always isolate")
• tests/test_checkpointing.py — all previously-red tests pass (test_gradient_matches_uncheckpointed, test_taylor_test_under_checkpointing, test_disk_gradient_matches_uncheckpointed, test_disk_taylor_test)
• Full suite: 86 passed, 1 xfailed (pre-existing, unrelated residual-timestepping defect)
• ruff check . and mypy src/dolfinx_adjoint clean

🤖 Generated with Claude Code