Revise the adjoint dependencies (#201)

* wip

* Revise the adjoint dependences

* falke8

Author: Ig-dolci

pyadjoint/checkpointing.py

@@ -406,11 +406,15 @@ def _(self, cp_action, progress_bar, markings, functional=None, **kwargs):
             current_step = self.tape.timesteps[step]
             for block in reversed(current_step):
                 block.evaluate_adj(markings=markings)
-                if not current_step._adj_deps_cleaned:
+                if not current_step._revised_adj_deps:
+                    # Update the adjoint dependency set.
+                    for deps in block.get_dependencies():
+                        if deps.marked_in_path:
+                            current_step.adjoint_dependencies.add(deps)
                     for out in block._outputs:
                         if not out.marked_in_path:
                             current_step.adjoint_dependencies.discard(out)
-            current_step._adj_deps_cleaned = True
+            current_step._revised_adj_deps = True
             # Output variables are used for the last time when running
             # backwards.
             to_keep = current_step.checkpointable_state

pyadjoint/tape.py

@@ -794,9 +794,10 @@ def __init__(self, blocks=()):
         # A dictionary mapping the block variables in the checkpointable state
         # to their checkpoint values.
         self._checkpoint = {}
-        # A flag to indicate whether the adjoint dependencies have been cleaned
-        # from the outputs not marked in the path.
-        self._adj_deps_cleaned = False
+        # Flag indicating whether the adjoint dependencies have been revised
+        # by removing outputs not marked in the path and adding checkpointable
+        # states that are marked in the path.
+        self._revised_adj_deps = False
 
     def copy(self, blocks=None):
         out = TimeStep(blocks or self)
@@ -826,8 +827,15 @@ def checkpoint(self, checkpointable_state, adj_dependencies, global_deps):
                         self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
 
             if adj_dependencies:
-                for var in self.adjoint_dependencies:
-                    self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
+                if self._revised_adj_deps:
+                    for var in self.adjoint_dependencies:
+                        self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
+                else:
+                    # The adjoint dependencies have not been revised yet. At this stage,
+                    # the block nodes are not marked in the path because the control variable(s)
+                    # are not yet determined.
+                    for var in self.adjoint_dependencies.union(self.checkpointable_state):
+                        self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
 
     def restore_from_checkpoint(self, from_storage):
         """Restore the block var checkpoints from the timestep checkpoint."""