From 7ecba7d3847143dac30a2332b5990a69c55548d4 Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Fri, 12 Jun 2026 12:11:38 +1000 Subject: [PATCH 1/3] Add regression test for SingleMemoryStorageSchedule long-range dependencies A block variable reused only every third timestep is absent from the immediately preceding step's adjoint_dependencies, so the checkpoint clearing for SingleMemoryStorageSchedule discarded it during forward replay and the reverse pass produced a wrong gradient (73765 instead of 3205 in this test). Companion to dolfin-adjoint/pyadjoint#248, which fixes the clearing condition; see dolfin-adjoint/pyadjoint#211 for the original report. --- .../adjoint/test_checkpointing_multistep.py | 51 ++++++++++++++++++- 1 file changed, 50 insertions(+), 1 deletion(-) diff --git a/tests/firedrake/adjoint/test_checkpointing_multistep.py b/tests/firedrake/adjoint/test_checkpointing_multistep.py index 61c80b9175..16450d1183 100644 --- a/tests/firedrake/adjoint/test_checkpointing_multistep.py +++ b/tests/firedrake/adjoint/test_checkpointing_multistep.py @@ -4,7 +4,8 @@ from firedrake.adjoint import * from .test_burgers_newton import _check_forward, \ _check_recompute, _check_reverse -from checkpoint_schedules import MixedCheckpointSchedule, StorageType +from checkpoint_schedules import MixedCheckpointSchedule, StorageType, \ + SingleMemoryStorageSchedule import numpy as np from collections import deque @@ -92,3 +93,51 @@ def test_validity(V): val_recomputed = J_hat(displacement_0) assert np.allclose(val_recomputed, val_recomputed0) assert np.allclose(dJ.dat.data_ro[:], dJ0.dat.data_ro[:]) + + +@pytest.mark.skipcomplex +def test_validity_single_memory_long_range(V): + """Long-range dependencies must survive SingleMemoryStorageSchedule. + + A variable that is reused only every third timestep is absent from the + immediately preceding step's adjoint dependencies, but its checkpoint + must not be cleared during the forward replay, otherwise the reverse + pass reconstructs a wrong value and the gradient is corrupted. See + https://github.com/dolfin-adjoint/pyadjoint/issues/211. + """ + def J_staggered(u_0): + tape = get_working_tape() + u = Function(V).assign(u_0) + r = Function(V) + for i in tape.timestepper(range(10)): + if i % 3 == 0: + # Refresh r only every third step: the resulting reuse gap + # is what the checkpoint clearing mishandled. Projecting r + # every step hides the bug. + r.project(1.01 * u) + u.project(r * u) + return assemble(u * u * dx) + + tape = get_working_tape() + tape.progress_bar = ProgressBar + u_0 = Function(V).assign(1.0) + # Without checkpointing. + val0 = J_staggered(u_0) + J_hat0 = ReducedFunctional(val0, Control(u_0)) + val_recomputed0 = J_hat0(u_0) + dJ0 = J_hat0.derivative() + tape.clear_tape() + + # With checkpointing. + tape.enable_checkpointing(SingleMemoryStorageSchedule()) + val = J_staggered(u_0) + J_hat = ReducedFunctional(val, Control(u_0)) + assert len(tape.timesteps) == 10 + # The functional must be re-evaluated *before* the derivative: the + # checkpoint clearing under test only runs during the forward replay + # triggered by this call. With derivative() first (as in test_validity + # above) the bug is not exercised. + val_recomputed = J_hat(u_0) + dJ = J_hat.derivative() + assert np.allclose(val_recomputed, val_recomputed0) + assert np.allclose(dJ.dat.data_ro[:], dJ0.dat.data_ro[:]) From 47afbe5f22fa7e8f1a32bca9db4e36189f917b34 Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Fri, 12 Jun 2026 12:11:38 +1000 Subject: [PATCH 2/3] TEMPORARY: test against pyadjoint PR #248 branch Drop this commit (restoring the pyadjoint-ad version pin) once dolfin-adjoint/pyadjoint#248 is merged and released. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index adb52bd8ed..e2beb12764 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ dependencies = [ "petsctools>=2026.0", "pkgconfig", "progress", - "pyadjoint-ad>=2026.4.0", + "pyadjoint-ad @ git+https://github.com/sghelichkhani/pyadjoint.git@sghelichkhani/singlemem-checkpoint-clearing", "pycparser", "pytools[siphash]", "requests", From 1f1fdaec1a43a4e3e9c7cd836fcd51957af5819d Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Fri, 12 Jun 2026 14:53:53 +1000 Subject: [PATCH 3/3] Add regression test for SingleMemoryStorageSchedule at a new control A variable that is never redefined keeps the live block variable of its Function, and the clearing for SingleMemoryStorageSchedule discarded its checkpoint at its last forward use even though the adjoint of that step still needs it as a linearisation point. The saved_output fallback to the live Function then supplies the taping-time value, corrupting the gradient whenever the functional is re-evaluated at a new control first (140.5536 instead of 187.4048 in this test). Companion to dolfin-adjoint/pyadjoint#248; see dolfin-adjoint/pyadjoint#260 for the report. --- .../adjoint/test_checkpointing_multistep.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/firedrake/adjoint/test_checkpointing_multistep.py b/tests/firedrake/adjoint/test_checkpointing_multistep.py index 16450d1183..a7eb3c90d3 100644 --- a/tests/firedrake/adjoint/test_checkpointing_multistep.py +++ b/tests/firedrake/adjoint/test_checkpointing_multistep.py @@ -141,3 +141,66 @@ def J_staggered(u_0): dJ = J_hat.derivative() assert np.allclose(val_recomputed, val_recomputed0) assert np.allclose(dJ.dat.data_ro[:], dJ0.dat.data_ro[:]) + + +@pytest.mark.skipcomplex +def test_validity_single_memory_new_control(V): + """Terminal variables must survive clearing at their last forward use. + + f is never redefined after step 0, so its block variable remains the + live one of the Function. If its checkpoint is cleared at its last + forward use, the adjoint of that step reads the linearisation point + through the saved_output fallback to the live Function, which still + holds the value from taping. The gradient is then wrong whenever the + functional has been re-evaluated at a new control first. See + https://github.com/dolfin-adjoint/pyadjoint/issues/260. + """ + def J_terminal(m): + tape = get_working_tape() + f = Function(V) + u = Function(V) + for i in tape.timestepper(iter(range(4))): + if i == 0: + # f is never redefined after this step. u must pick up + # control dependence here: the stale linearisation point + # weights the derivative with respect to u, so without + # this the wrong weight multiplies a path carrying no + # control sensitivity and the gradient is accidentally + # right. + f.project(2 * m) + u.project(m) + elif i == 2: + # Last forward use of f. The adjoint of this projection + # needs f's value as the linearisation point for the + # derivative with respect to u. + u.project(u * f) + else: + u.project(Constant(1.1) * u) + return assemble(u * u * dx) + + tape = get_working_tape() + m = Function(V).assign(1.0) + controls = [Function(V).assign(2.0), Function(V).assign(3.0)] + + # Without checkpointing. + J_hat0 = ReducedFunctional(J_terminal(m), Control(m)) + vals0 = [] + grads0 = [] + for m_new in controls: + vals0.append(J_hat0(m_new)) + grads0.append(J_hat0.derivative().dat.data_ro.copy()) + tape.clear_tape() + + # With checkpointing. + tape.enable_checkpointing(SingleMemoryStorageSchedule()) + J_hat = ReducedFunctional(J_terminal(m), Control(m)) + for m_new, val0, grad0 in zip(controls, vals0, grads0): + # Each cycle re-evaluates at a new control *before* the + # derivative: the stale fallback only shows away from the taping + # point. The second cycle exercises the clearing that runs once + # the adjoint dependencies have been revised by the first + # reverse pass. + val = J_hat(m_new) + dJ = J_hat.derivative() + assert np.allclose(val, val0) + assert np.allclose(dJ.dat.data_ro[:], grad0)