From a1e40b7a06353163bed19419aceb34108edb81db Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Sun, 15 Feb 2026 22:48:16 +1100 Subject: [PATCH 1/3] Fix SingleMemoryStorageSchedule incorrectly clearing checkpoints Fixes #211. --- pyadjoint/checkpointing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyadjoint/checkpointing.py b/pyadjoint/checkpointing.py index a23c2d95..c20e1b49 100644 --- a/pyadjoint/checkpointing.py +++ b/pyadjoint/checkpointing.py @@ -361,7 +361,7 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs): # Handle the case for SingleMemoryStorageSchedule if isinstance(self._schedule, SingleMemoryStorageSchedule): - if step > 1 and var not in self.tape.timesteps[step - 1].adjoint_dependencies: + if var.output.block_variable is var: var.checkpoint = None continue From be651ab19f559b8bca3cd2e5d1f1efd484bddce4 Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Fri, 12 Jun 2026 13:59:33 +1000 Subject: [PATCH 2/3] Only clear SingleMemoryStorageSchedule checkpoints once adjoint dependencies are revised A variable visited by the clear-down loop at its last forward use is by construction a dependency of a block in that step, so the adjoint of the step may still need it as a linearisation point. Clearing it and relying on saved_output falling back to the live function returns the stale taped value whenever the functional has been re-evaluated at a new control. Keep everything until a reverse pass has revised the step's adjoint dependencies, then clear only what is provably not needed. Fixes #260. --- pyadjoint/checkpointing.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pyadjoint/checkpointing.py b/pyadjoint/checkpointing.py index c20e1b49..848adffb 100644 --- a/pyadjoint/checkpointing.py +++ b/pyadjoint/checkpointing.py @@ -361,7 +361,16 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs): # Handle the case for SingleMemoryStorageSchedule if isinstance(self._schedule, SingleMemoryStorageSchedule): - if var.output.block_variable is var: + # `var` is used by a block in this step, so the adjoint of + # this step may still need it as a linearisation point. + # Only clear once the adjoint dependencies have been + # revised by a reverse pass and `var` is provably not one + # of them; before that, keeping every dependency in memory + # is exactly what this schedule promises. + if ( + current_step._revised_adj_deps + and var not in current_step.adjoint_dependencies + ): var.checkpoint = None continue From 4a9d3b6fbb3b0802e9edd5e0aecc7e3064ae315a Mon Sep 17 00:00:00 2001 From: Sia Ghelichkhan Date: Tue, 16 Jun 2026 13:39:07 +1000 Subject: [PATCH 3/3] Note the forward-only memory trade-off in the SingleMemory clearing comment --- pyadjoint/checkpointing.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyadjoint/checkpointing.py b/pyadjoint/checkpointing.py index 848adffb..36f9fb0c 100644 --- a/pyadjoint/checkpointing.py +++ b/pyadjoint/checkpointing.py @@ -366,7 +366,10 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs): # Only clear once the adjoint dependencies have been # revised by a reverse pass and `var` is provably not one # of them; before that, keeping every dependency in memory - # is exactly what this schedule promises. + # is exactly what this schedule promises. The trade-off is + # that a forward-only recomputation retains the conservative + # dependency set and so holds taping-time memory; the more + # precise clearing only takes effect after the first reverse. if ( current_step._revised_adj_deps and var not in current_step.adjoint_dependencies