From 9e84696bd76ca1c27bd030c86ae32df13307e4a3 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 12:27:01 +0100
Subject: [PATCH 01/29] Reverse-over-forward AD: AssembleBlock

---
 firedrake/adjoint/__init__.py                 |  4 +-
 firedrake/adjoint_utils/blocks/assembly.py    | 28 +++++++-
 firedrake/adjoint_utils/blocks/block_utils.py | 64 +++++++++++++++++++
 .../test_adjoint_reverse_over_forward.py      | 39 +++++++++++
 4 files changed, 133 insertions(+), 2 deletions(-)
 create mode 100644 tests/regression/test_adjoint_reverse_over_forward.py

diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
index 8373a4285a..c3f4c8a00a 100644
--- a/firedrake/adjoint/__init__.py
+++ b/firedrake/adjoint/__init__.py
@@ -19,7 +19,9 @@
 
 from pyadjoint.tape import Tape, set_working_tape, get_working_tape, \
                             pause_annotation, continue_annotation, \
-                            stop_annotating, annotate_tape  # noqa F401
+                            stop_annotating, annotate_tape, \
+                            pause_reverse_over_forward, continue_reverse_over_forward, \
+                            stop_reverse_over_forward  # noqa F401
 from pyadjoint.reduced_functional import ReducedFunctional  # noqa F401
 from firedrake.adjoint_utils.checkpointing import \
     enable_disk_checkpointing, pause_disk_checkpointing, \
diff --git a/firedrake/adjoint_utils/blocks/assembly.py b/firedrake/adjoint_utils/blocks/assembly.py
index 8ba790ec74..8d199725ba 100644
--- a/firedrake/adjoint_utils/blocks/assembly.py
+++ b/firedrake/adjoint_utils/blocks/assembly.py
@@ -4,7 +4,7 @@
 from ufl.formatting.ufl2unicode import ufl2unicode
 from pyadjoint import Block, AdjFloat, create_overloaded_object
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
-from .block_utils import isconstant
+from .block_utils import isconstant, restored_outputs
 
 
 class AssembleBlock(Block):
@@ -145,6 +145,32 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             dform = firedrake.assemble(dform)
         return dform
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        form = self.form
+
+        tlm_rhs = 0
+        for block_variable in self.get_dependencies():
+            dep = block_variable.output
+            tlm_dep = block_variable.tlm_value
+            if tlm_dep is not None:
+                if isinstance(dep, firedrake.MeshGeometry):
+                    dep = firedrake.SpatialCoordinate(dep)
+                    tlm_rhs = tlm_rhs + firedrake.derivative(
+                        form, dep, tlm_dep)
+                else:
+                    tlm_rhs = tlm_rhs + firedrake.action(
+                        firedrake.derivative(form, dep), tlm_dep)
+
+        x.tlm_value = None
+        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+            return
+        tau_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if tau_rhs.empty():
+            return
+        with restored_outputs(x, restore=lambda x: x in form.coefficients()):
+            x.tlm_value = firedrake.assemble(tau_rhs)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, adj_inputs,
diff --git a/firedrake/adjoint_utils/blocks/block_utils.py b/firedrake/adjoint_utils/blocks/block_utils.py
index ee3b95ff20..55fa070c7f 100644
--- a/firedrake/adjoint_utils/blocks/block_utils.py
+++ b/firedrake/adjoint_utils/blocks/block_utils.py
@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 import firedrake
 
 
@@ -8,3 +10,65 @@ def isconstant(expr):
     if isinstance(expr, firedrake.Constant):
         raise ValueError("Firedrake Constant requires a domain to work with pyadjoint")
     return isinstance(expr, firedrake.Function) and expr.ufl_element().family() == "Real"
+
+
+@contextmanager
+def restored_outputs(*X, restore=None):
+    """Construct a context manager which can be used to temporarily restore
+    block variable outputs to saved values.
+
+    Parameters
+    ----------
+    X : tuple[BlockVariable]
+        Block variables to temporarily restore.
+    restore : callable
+        Can be used to exclude variables. Only inputs for which
+        `restore(x.output)` is true have their outputs temporarily restored.
+
+    Returns
+    -------
+
+    The context manager.
+
+    Notes
+    -----
+
+    A forward operation is allowed to modify the original variable, e.g. in
+
+    .. code-block:: python3
+
+        solve(inner(trial, test) * dx
+              == inner(x * x, test) * dx,
+              x)
+
+    `x` has two versions: the input and the output. Reverse-over-forward AD
+    requires that we use the symbolic representation `x`, but with input value
+    `x.block_variable.saved_output`. A context manager can be used to
+    temporarily restore the value of `x` so that we can perform and annotate
+    a tangent-linear operation,
+
+    .. code-block:: python3
+
+        with restored_outputs(x):
+            # The value of x is now x.block_variable.saved_output
+            solve(inner(trial, test) * dx
+                  == 2 * inner(x * x.block_variable.tlm_value, test) * dx,
+                  x.block_variable.tlm_value)
+        # The value of x is again the output from the forward solve(...)
+    """
+
+    if restore is None:
+        def restore(x):
+            return True
+
+    X = tuple(x for x in X if restore(x.output))
+    X_old = tuple(x.output._ad_copy(x) for x in X)
+    for x in X:
+        # Ideally would use a generic _ad_assign here
+        x.output.assign(x.output.block_variable.saved_output)
+    try:
+        yield
+    finally:
+        for x, x_old in zip(X, X_old):
+            # Ideally would use a generic _ad_assign here
+            x.output.assign(x_old)
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
new file mode 100644
index 0000000000..b39da1a7d3
--- /dev/null
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -0,0 +1,39 @@
+import numpy as np
+import pytest
+
+from firedrake import *
+from firedrake.adjoint import *
+from firedrake.__future__ import *
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _():
+    get_working_tape().clear_tape()
+    pause_annotation()
+    pause_reverse_over_forward()
+    yield
+    get_working_tape().clear_tape()
+    pause_annotation()
+    pause_reverse_over_forward()
+
+
+def test_assembly():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(Constant(1.0))
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+    continue_annotation()
+    continue_reverse_over_forward()
+    J = assemble(u * u * dx)
+    pause_annotation()
+    pause_reverse_over_forward()
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)

From bbbe8419e15c7c3a15c8bfb959759ebd75d722a5 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 13:07:15 +0100
Subject: [PATCH 02/29] In-place assignment

---
 firedrake/adjoint_utils/constant.py | 3 +++
 firedrake/adjoint_utils/function.py | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/firedrake/adjoint_utils/constant.py b/firedrake/adjoint_utils/constant.py
index fed5978a6d..c9606981d4 100644
--- a/firedrake/adjoint_utils/constant.py
+++ b/firedrake/adjoint_utils/constant.py
@@ -104,6 +104,9 @@ def _ad_assign_numpy(dst, src, offset):
     def _ad_to_list(m):
         return m.dat.data_ro.reshape(-1).tolist()
 
+    def _ad_assign(self, other):
+        self.assign(other)
+
     def _ad_copy(self):
         return self._constant_from_values()
 
diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py
index d56438fec7..5492ad04a4 100644
--- a/firedrake/adjoint_utils/function.py
+++ b/firedrake/adjoint_utils/function.py
@@ -352,6 +352,9 @@ def _ad_to_list(m):
 
         return m_a.tolist()
 
+    def _ad_assign(self, other):
+        self.assign(other)
+
     def _ad_copy(self):
         from firedrake import Function
 

From 18a50008b98accb8a9c2dd020b47ac4e73b79187 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 13:08:17 +0100
Subject: [PATCH 03/29] Remove restored_output, now handled within pyadjoint

---
 firedrake/adjoint_utils/blocks/assembly.py    |  5 +-
 firedrake/adjoint_utils/blocks/block_utils.py | 64 -------------------
 2 files changed, 2 insertions(+), 67 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/assembly.py b/firedrake/adjoint_utils/blocks/assembly.py
index 8d199725ba..4f53465b29 100644
--- a/firedrake/adjoint_utils/blocks/assembly.py
+++ b/firedrake/adjoint_utils/blocks/assembly.py
@@ -4,7 +4,7 @@
 from ufl.formatting.ufl2unicode import ufl2unicode
 from pyadjoint import Block, AdjFloat, create_overloaded_object
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
-from .block_utils import isconstant, restored_outputs
+from .block_utils import isconstant
 
 
 class AssembleBlock(Block):
@@ -168,8 +168,7 @@ def solve_tlm(self):
         tau_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
         if tau_rhs.empty():
             return
-        with restored_outputs(x, restore=lambda x: x in form.coefficients()):
-            x.tlm_value = firedrake.assemble(tau_rhs)
+        x.tlm_value = firedrake.assemble(tau_rhs)
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
diff --git a/firedrake/adjoint_utils/blocks/block_utils.py b/firedrake/adjoint_utils/blocks/block_utils.py
index 55fa070c7f..ee3b95ff20 100644
--- a/firedrake/adjoint_utils/blocks/block_utils.py
+++ b/firedrake/adjoint_utils/blocks/block_utils.py
@@ -1,5 +1,3 @@
-from contextlib import contextmanager
-
 import firedrake
 
 
@@ -10,65 +8,3 @@ def isconstant(expr):
     if isinstance(expr, firedrake.Constant):
         raise ValueError("Firedrake Constant requires a domain to work with pyadjoint")
     return isinstance(expr, firedrake.Function) and expr.ufl_element().family() == "Real"
-
-
-@contextmanager
-def restored_outputs(*X, restore=None):
-    """Construct a context manager which can be used to temporarily restore
-    block variable outputs to saved values.
-
-    Parameters
-    ----------
-    X : tuple[BlockVariable]
-        Block variables to temporarily restore.
-    restore : callable
-        Can be used to exclude variables. Only inputs for which
-        `restore(x.output)` is true have their outputs temporarily restored.
-
-    Returns
-    -------
-
-    The context manager.
-
-    Notes
-    -----
-
-    A forward operation is allowed to modify the original variable, e.g. in
-
-    .. code-block:: python3
-
-        solve(inner(trial, test) * dx
-              == inner(x * x, test) * dx,
-              x)
-
-    `x` has two versions: the input and the output. Reverse-over-forward AD
-    requires that we use the symbolic representation `x`, but with input value
-    `x.block_variable.saved_output`. A context manager can be used to
-    temporarily restore the value of `x` so that we can perform and annotate
-    a tangent-linear operation,
-
-    .. code-block:: python3
-
-        with restored_outputs(x):
-            # The value of x is now x.block_variable.saved_output
-            solve(inner(trial, test) * dx
-                  == 2 * inner(x * x.block_variable.tlm_value, test) * dx,
-                  x.block_variable.tlm_value)
-        # The value of x is again the output from the forward solve(...)
-    """
-
-    if restore is None:
-        def restore(x):
-            return True
-
-    X = tuple(x for x in X if restore(x.output))
-    X_old = tuple(x.output._ad_copy(x) for x in X)
-    for x in X:
-        # Ideally would use a generic _ad_assign here
-        x.output.assign(x.output.block_variable.saved_output)
-    try:
-        yield
-    finally:
-        for x, x_old in zip(X, X_old):
-            # Ideally would use a generic _ad_assign here
-            x.output.assign(x_old)

From 73c30518d8f5db20de0e090a904470cd5c15d07b Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 14:04:01 +0100
Subject: [PATCH 04/29] Reverse-over-forward AD: ConstantAssignBlock

---
 firedrake/adjoint_utils/blocks/constant.py    | 19 ++++++++--
 .../test_adjoint_reverse_over_forward.py      | 36 +++++++++++++++----
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/constant.py b/firedrake/adjoint_utils/blocks/constant.py
index d5e580428a..7b38a063fd 100644
--- a/firedrake/adjoint_utils/blocks/constant.py
+++ b/firedrake/adjoint_utils/blocks/constant.py
@@ -1,6 +1,6 @@
-from pyadjoint import Block, OverloadedType
+import firedrake
 import numpy
-
+from pyadjoint import Block, OverloadedType
 from pyadjoint.reduced_functional_numpy import gather
 from .block_utils import isconstant
 
@@ -70,6 +70,21 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             values = values.values()
         return constant_from_values(block_variable.output, values)
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        if len(x.output.ufl_shape) == 0:
+            x.tlm_value = firedrake.Constant(0.0)
+        else:
+            x.tlm_value = firedrake.Constant(
+                numpy.reshape(numpy.zeros_like(x.output.values()), x.output.ufl_shape))
+        if self.assigned_list:
+            # Not reachable?
+            raise NotImplementedError
+        else:
+            dep, = self.get_dependencies()
+            if dep.tlm_value is not None:
+                x.tlm_value.assign(dep.tlm_value)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, hessian_inputs,
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index b39da1a7d3..cf2d183a7f 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import numpy as np
 import pytest
 
@@ -6,7 +7,7 @@
 from firedrake.__future__ import *
 
 
-@pytest.fixture(autouse=True, scope="module")
+@pytest.fixture(autouse=True)
 def _():
     get_working_tape().clear_tape()
     pause_annotation()
@@ -17,6 +18,16 @@ def _():
     pause_reverse_over_forward()
 
 
+@contextmanager
+def reverse_over_forward():
+    continue_annotation()
+    continue_reverse_over_forward()
+    yield
+    pause_annotation()
+    pause_reverse_over_forward()
+
+
+@pytest.mark.skipcomplex
 def test_assembly():
     mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
@@ -27,13 +38,26 @@ def test_assembly():
     zeta = Function(space, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
-    continue_annotation()
-    continue_reverse_over_forward()
-    J = assemble(u * u * dx)
-    pause_annotation()
-    pause_reverse_over_forward()
+    with reverse_over_forward():
+        J = assemble(u * u * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
                        assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_constant_assignment():
+    a = Constant(2.5)
+    a.block_variable.tlm_value = Constant(-2.0)
+
+    with reverse_over_forward():
+        b = Constant(0.0).assign(a)
+
+    assert float(b.block_variable.tlm_value) == -2.0
+
+    # Minimal test that the TLM operations are on the tape
+    _ = compute_gradient(b.block_variable.tlm_value, Control(a.block_variable.tlm_value))
+    adj_value = a.block_variable.tlm_value.block_variable.adj_value
+    assert float(adj_value) == 1.0

From ceab69cb83607f1487ce0ccd13fa16deeec88179 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 14:18:11 +0100
Subject: [PATCH 05/29] Rename variables

---
 firedrake/adjoint_utils/blocks/assembly.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/assembly.py b/firedrake/adjoint_utils/blocks/assembly.py
index 4f53465b29..05b7672300 100644
--- a/firedrake/adjoint_utils/blocks/assembly.py
+++ b/firedrake/adjoint_utils/blocks/assembly.py
@@ -165,10 +165,10 @@ def solve_tlm(self):
         x.tlm_value = None
         if isinstance(tlm_rhs, int) and tlm_rhs == 0:
             return
-        tau_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
-        if tau_rhs.empty():
+        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if tlm_rhs.empty():
             return
-        x.tlm_value = firedrake.assemble(tau_rhs)
+        x.tlm_value = firedrake.assemble(tlm_rhs)
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):

From e42dd5e72d5c02ad25c2edff95e68b18d006bc9a Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 14:30:17 +0100
Subject: [PATCH 06/29] Reverse-over-forward AD: FunctionAssignBlock. Minor
 test edits.

---
 firedrake/adjoint_utils/blocks/function.py    | 19 ++++++++++++
 .../test_adjoint_reverse_over_forward.py      | 29 ++++++++++++++++---
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py
index fc5be8486a..8146c6955b 100644
--- a/firedrake/adjoint_utils/blocks/function.py
+++ b/firedrake/adjoint_utils/blocks/function.py
@@ -129,6 +129,25 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
 
         return dudm
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        expr = self.expr
+
+        tlm_rhs = 0
+        for block_variable in self.get_dependencies():
+            dep = block_variable.output
+            tlm_dep = block_variable.tlm_value
+            if tlm_dep is not None:
+                tlm_rhs = tlm_rhs + ufl.derivative(expr, dep, tlm_dep)
+                
+        x.tlm_value = None
+        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+            return
+        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if isinstance(tlm_rhs, ufl.constantvalue.Zero):
+            return
+        x.tlm_value = firedrake.Function(x.output.function_space()).assign(tlm_rhs)
+
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
         return self.prepare_evaluate_adj(inputs, hessian_inputs,
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index cf2d183a7f..bd71c366f8 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -34,17 +34,17 @@ def test_assembly():
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
 
-    u = Function(space, name="u").interpolate(Constant(1.0))
+    u = Function(space, name="u").interpolate(X[0])
     zeta = Function(space, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
     with reverse_over_forward():
-        J = assemble(u * u * dx)
+        J = assemble(u * u.dx(0) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+                       assemble(inner(zeta.dx(0), test) * dx + inner(zeta, test.dx(0)) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -57,7 +57,28 @@ def test_constant_assignment():
 
     assert float(b.block_variable.tlm_value) == -2.0
 
-    # Minimal test that the TLM operations are on the tape
+    # Minimal test that the TLM operation is on the tape
     _ = compute_gradient(b.block_variable.tlm_value, Control(a.block_variable.tlm_value))
     adj_value = a.block_variable.tlm_value.block_variable.adj_value
     assert float(adj_value) == 1.0
+
+
+@pytest.mark.skipcomplex
+def test_function_assignment():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+    with reverse_over_forward():
+        v = Function(space, name="v").assign(-3 * u)
+        J = assemble(v * v.dx(0) * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(9 * inner(zeta.dx(0), test) * dx + 9 * inner(zeta, test.dx(0)) * dx).dat.data_ro)

From 26865be08d65deaf816c451f6e40e0999ae6ad8d Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 15:06:44 +0100
Subject: [PATCH 07/29] Reverse-over-forward AD: GenericSolveBlock

---
 firedrake/adjoint_utils/blocks/solving.py     | 44 +++++++++++++++++++
 .../test_adjoint_reverse_over_forward.py      | 23 ++++++++++
 2 files changed, 67 insertions(+)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 7a57e883c3..05b7ac88f5 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -317,6 +317,50 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             dFdm, dudm, bcs
         )
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        if self.linear:
+            form = firedrake.action(self.lhs, x.output) - self.rhs
+        else:
+            form = self.lhs
+
+        tlm_rhs = 0
+        tlm_bcs = []
+        for block_variable in self.get_dependencies():            
+            dep = block_variable.output
+            if dep == x.output:
+                continue
+            tlm_dep = block_variable.tlm_value
+            if isinstance(dep, firedrake.DirichletBC):
+                if tlm_value is None:
+                    tlm_bcs.append(dep.reconstruct(g=0))
+                else:
+                    tlm_bcs.append(tlm_value)
+            elif tlm_dep is not None:
+                if isinstance(dep, firedrake.MeshGeometry):
+                    dep = firedrake.SpatialCoordinate(dep)
+                    tlm_rhs = tlm_rhs - firedrake.derivative(
+                        form, dep, tlm_dep)
+                else:
+                    tlm_rhs = tlm_rhs - firedrake.action(
+                        firedrake.derivative(form, dep), tlm_dep)
+
+        x.tlm_value = None
+        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+            return
+        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+        if tlm_rhs.empty():
+            return
+
+        if self.linear:
+            J = self.lhs
+        else:
+            J = firedrake.derivative(form, x.output, firedrake.TrialFunction(x.output.function_space()))
+
+        x.tlm_value = firedrake.Function(x.output.function_space())
+        firedrake.solve(J == tlm_rhs, x.tlm_value, tlm_bcs, *self.forward_args,
+                        **self.forward_kwargs)
+
     def _assemble_and_solve_tlm_eq(self, dFdu, dFdm, dudm, bcs):
         return self._assembled_solve(dFdu, dFdm, dudm, bcs)
 
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index bd71c366f8..aac8a037d4 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -82,3 +82,26 @@ def test_function_assignment():
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
                        assemble(9 * inner(zeta.dx(0), test) * dx + 9 * inner(zeta, test.dx(0)) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_project():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+    space_0 = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
+
+    with reverse_over_forward():
+        v = Function(space_0, name="v").project(u)
+        J = assemble(v * v * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(2 * inner(Function(space_0).project(zeta), test) * dx).dat.data_ro)

From b9fe7322988264e17ed667157014f18f0f0494cd Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 15:32:47 +0100
Subject: [PATCH 08/29] Reverse-over-forward AD: SupermeshProjectBlock

---
 firedrake/adjoint_utils/blocks/solving.py     | 10 +++++
 .../test_adjoint_reverse_over_forward.py      | 38 +++++++++++++++----
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 05b7ac88f5..d435237990 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -842,6 +842,7 @@ def __init__(self, source, target_space, target, bcs=[], **kwargs):
             mesh = target_space.mesh()
         self.source_space = source.function_space()
         self.target_space = target_space
+        self._kwargs = dict(kwargs)
         self.projector = firedrake.Projector(source, target_space, **kwargs)
 
         # Assemble mixed mass matrix
@@ -922,6 +923,15 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
                                              prepared)
         return dJdm
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        dep, = self.get_dependencies()
+        if dep.tlm_value is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = firedrake.Function(x.output.function_space())
+            firedrake.project(dep.tlm_value, x.tlm_value, **self._kwargs)
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
                                    relevant_dependencies, prepared=None):
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index aac8a037d4..93a2e0ccea 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -86,22 +86,44 @@ def test_function_assignment():
 
 @pytest.mark.skipcomplex
 def test_project():
-    mesh = UnitIntervalMesh(10)
+    mesh = UnitSquareMesh(10, 10)
     X = SpatialCoordinate(mesh)
-    space = FunctionSpace(mesh, "Lagrange", 1)
-    test = TestFunction(space)
+    space_a = FunctionSpace(mesh, "Lagrange", 1)
+    space_b = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
+    test_a = TestFunction(space_a)
 
-    u = Function(space, name="u").interpolate(X[0] - 0.5)
-    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space_a, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
-    space_0 = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
+    with reverse_over_forward():
+        v = Function(space_b, name="v").project(u)
+        J = assemble(v * v * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(2 * inner(Function(space_b).project(zeta), test_a) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_supermesh_project():
+    mesh_a = UnitSquareMesh(10, 10)
+    mesh_b = UnitSquareMesh(5, 20)
+    X_a = SpatialCoordinate(mesh_a)
+    space_a = FunctionSpace(mesh_a, "Lagrange", 1)
+    space_b = FunctionSpace(mesh_b, "Discontinuous Lagrange", 0)
+    test_a = TestFunction(space_a)
+
+    u = Function(space_a, name="u").interpolate(X_a[0] - 0.5)
+    zeta = Function(space_a, name="tlm_u").interpolate(X_a[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
     with reverse_over_forward():
-        v = Function(space_0, name="v").project(u)
+        v = Function(space_b, name="v").project(u)
         J = assemble(v * v * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(Function(space_0).project(zeta), test) * dx).dat.data_ro)
+                       assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)

From 7dbd38022c06f86727ed46c793bfa05a5506e073 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 16:09:06 +0100
Subject: [PATCH 09/29] Fix

---
 firedrake/adjoint_utils/blocks/solving.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index d435237990..3790a83860 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -332,10 +332,10 @@ def solve_tlm(self):
                 continue
             tlm_dep = block_variable.tlm_value
             if isinstance(dep, firedrake.DirichletBC):
-                if tlm_value is None:
+                if tlm_dep is None:
                     tlm_bcs.append(dep.reconstruct(g=0))
                 else:
-                    tlm_bcs.append(tlm_value)
+                    tlm_bcs.append(tlm_dep)
             elif tlm_dep is not None:
                 if isinstance(dep, firedrake.MeshGeometry):
                     dep = firedrake.SpatialCoordinate(dep)

From 4c74bd18576af187d38646d5ac1ed3841606b4ae Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 16:09:50 +0100
Subject: [PATCH 10/29] flake8

---
 firedrake/adjoint_utils/blocks/function.py | 2 +-
 firedrake/adjoint_utils/blocks/solving.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py
index 8146c6955b..2bce0eef64 100644
--- a/firedrake/adjoint_utils/blocks/function.py
+++ b/firedrake/adjoint_utils/blocks/function.py
@@ -139,7 +139,7 @@ def solve_tlm(self):
             tlm_dep = block_variable.tlm_value
             if tlm_dep is not None:
                 tlm_rhs = tlm_rhs + ufl.derivative(expr, dep, tlm_dep)
-                
+
         x.tlm_value = None
         if isinstance(tlm_rhs, int) and tlm_rhs == 0:
             return
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 3790a83860..2ca1c86c92 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -326,7 +326,7 @@ def solve_tlm(self):
 
         tlm_rhs = 0
         tlm_bcs = []
-        for block_variable in self.get_dependencies():            
+        for block_variable in self.get_dependencies():
             dep = block_variable.output
             if dep == x.output:
                 continue

From f10d7bbe4caa9fe1513a66651642561b8fdf2ae9 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 18:19:37 +0100
Subject: [PATCH 11/29] Reverse-over-forward AD: FunctionAssignBlock,
 self.expr==None case

---
 firedrake/adjoint_utils/blocks/function.py    | 35 +++++++++++--------
 .../test_adjoint_reverse_over_forward.py      | 29 ++++++++++++---
 2 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py
index 2bce0eef64..1c05ad825d 100644
--- a/firedrake/adjoint_utils/blocks/function.py
+++ b/firedrake/adjoint_utils/blocks/function.py
@@ -133,20 +133,27 @@ def solve_tlm(self):
         x, = self.get_outputs()
         expr = self.expr
 
-        tlm_rhs = 0
-        for block_variable in self.get_dependencies():
-            dep = block_variable.output
-            tlm_dep = block_variable.tlm_value
-            if tlm_dep is not None:
-                tlm_rhs = tlm_rhs + ufl.derivative(expr, dep, tlm_dep)
-
-        x.tlm_value = None
-        if isinstance(tlm_rhs, int) and tlm_rhs == 0:
-            return
-        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
-        if isinstance(tlm_rhs, ufl.constantvalue.Zero):
-            return
-        x.tlm_value = firedrake.Function(x.output.function_space()).assign(tlm_rhs)
+        if expr is None:
+            other, = self.get_dependencies()
+            if other.tlm_value is None:
+                x.tlm_value = None
+            else:
+                x.tlm_value = firedrake.Function(x.output.function_space()).assign(other.tlm_value)
+        else:
+            tlm_rhs = 0
+            for block_variable in self.get_dependencies():
+                dep = block_variable.output
+                tlm_dep = block_variable.tlm_value
+                if tlm_dep is not None:
+                    tlm_rhs = tlm_rhs + ufl.derivative(expr, dep, tlm_dep)
+
+            x.tlm_value = None
+            if isinstance(tlm_rhs, int) and tlm_rhs == 0:
+                return
+            tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+            if isinstance(tlm_rhs, ufl.constantvalue.Zero):
+                return
+            x.tlm_value = firedrake.Function(x.output.function_space()).assign(tlm_rhs)
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 93a2e0ccea..85ebcb9b19 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -39,12 +39,12 @@ def test_assembly():
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
     with reverse_over_forward():
-        J = assemble(u * u.dx(0) * dx)
+        J = assemble(u * u * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
-                       assemble(inner(zeta.dx(0), test) * dx + inner(zeta, test.dx(0)) * dx).dat.data_ro)
+                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -74,14 +74,35 @@ def test_function_assignment():
     zeta = Function(space, name="tlm_u").interpolate(X[0])
     u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
+    with reverse_over_forward():
+        v = Function(space, name="v").assign(u)
+        J = assemble(v * v * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+
+
+@pytest.mark.skipcomplex
+def test_function_assignment_expr():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    u = Function(space, name="u").interpolate(X[0] - 0.5)
+    zeta = Function(space, name="tlm_u").interpolate(X[0])
+    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
     with reverse_over_forward():
         v = Function(space, name="v").assign(-3 * u)
-        J = assemble(v * v.dx(0) * dx)
+        J = assemble(v * v * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
-                       assemble(9 * inner(zeta.dx(0), test) * dx + 9 * inner(zeta, test.dx(0)) * dx).dat.data_ro)
+                       assemble(18 * inner(zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex

From a264a879cbd1a262adaa1a0d7f884c0b6ae8c64a Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 18:39:24 +0100
Subject: [PATCH 12/29] Reverse-over-forward AD: SubfunctionBlock and
 FunctionMergeBlock

---
 firedrake/adjoint_utils/blocks/function.py    | 20 +++++++++++++++
 .../test_adjoint_reverse_over_forward.py      | 25 +++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py
index 1c05ad825d..30c4941448 100644
--- a/firedrake/adjoint_utils/blocks/function.py
+++ b/firedrake/adjoint_utils/blocks/function.py
@@ -237,6 +237,15 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
                                prepared=None):
         return firedrake.Function.sub(tlm_inputs[0], self.idx)
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        dep, = self.get_dependencies()
+        tlm_dep = dep.tlm_value
+        if tlm_dep is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = firedrake.Function(x.output.function_space()).assign(tlm_dep.sub(self.idx))
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
                                    relevant_dependencies, prepared=None):
@@ -279,6 +288,17 @@ def evaluate_tlm(self):
             type(output.output).assign(f.sub(self.idx), tlm_input)
         )
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        tlm_dep = self.get_dependencies()[0].tlm_value
+        if tlm_dep is None:
+            if x.tlm_value is not None:
+                x.tlm_value.sub(self.idx).zero()
+        else:
+            if x.tlm_value is None:
+                x.tlm_value = type(x.output)(x.output.function_space())
+            x.tlm_value.sub(self.idx).assign(tlm_dep)
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx,
                                    relevant_dependencies, prepared=None):
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 85ebcb9b19..80ca2e0e12 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -105,6 +105,31 @@ def test_function_assignment_expr():
                        assemble(18 * inner(zeta, test) * dx).dat.data_ro)
 
 
+@pytest.mark.skipcomplex
+@pytest.mark.parametrize("idx", [0, 1])
+def test_subfunction(idx):
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1) * FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    with reverse_over_forward():
+        u = Function(space, name="u")
+        u.sub(idx).interpolate(-2 * X[0])
+        zeta = Function(space, name="zeta")
+        zeta.sub(idx).interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+        v = Function(space, name="v")
+        v.sub(idx).assign(u.sub(idx))
+        J = assemble(u.sub(idx) * v.sub(idx) * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.sub(idx).dat.data_ro,
+                       assemble(2 * inner(zeta[idx], test[idx]) * dx).dat.data_ro[idx])
+
+
 @pytest.mark.skipcomplex
 def test_project():
     mesh = UnitSquareMesh(10, 10)

From 67746883e85780292df9ba7295a57b52041ddb71 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 18:42:41 +0100
Subject: [PATCH 13/29] Expand annotation in tests

---
 .../test_adjoint_reverse_over_forward.py      | 46 +++++++++----------
 1 file changed, 23 insertions(+), 23 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 80ca2e0e12..2cada1e4a2 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -34,11 +34,11 @@ def test_assembly():
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
 
-    u = Function(space, name="u").interpolate(X[0])
-    zeta = Function(space, name="tlm_u").interpolate(X[0])
-    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
-
     with reverse_over_forward():
+        u = Function(space, name="u").interpolate(X[0])
+        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
         J = assemble(u * u * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
@@ -49,10 +49,10 @@ def test_assembly():
 
 @pytest.mark.skipcomplex
 def test_constant_assignment():
-    a = Constant(2.5)
-    a.block_variable.tlm_value = Constant(-2.0)
-
     with reverse_over_forward():
+        a = Constant(2.5)
+        a.block_variable.tlm_value = Constant(-2.0)
+
         b = Constant(0.0).assign(a)
 
     assert float(b.block_variable.tlm_value) == -2.0
@@ -70,11 +70,11 @@ def test_function_assignment():
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
 
-    u = Function(space, name="u").interpolate(X[0] - 0.5)
-    zeta = Function(space, name="tlm_u").interpolate(X[0])
-    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
-
     with reverse_over_forward():
+        u = Function(space, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
         v = Function(space, name="v").assign(u)
         J = assemble(v * v * dx)
 
@@ -91,11 +91,11 @@ def test_function_assignment_expr():
     space = FunctionSpace(mesh, "Lagrange", 1)
     test = TestFunction(space)
 
-    u = Function(space, name="u").interpolate(X[0] - 0.5)
-    zeta = Function(space, name="tlm_u").interpolate(X[0])
-    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
-
     with reverse_over_forward():
+        u = Function(space, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
         v = Function(space, name="v").assign(-3 * u)
         J = assemble(v * v * dx)
 
@@ -138,11 +138,11 @@ def test_project():
     space_b = FunctionSpace(mesh, "Discontinuous Lagrange", 0)
     test_a = TestFunction(space_a)
 
-    u = Function(space_a, name="u").interpolate(X[0] - 0.5)
-    zeta = Function(space_a, name="tlm_u").interpolate(X[0])
-    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
-
     with reverse_over_forward():
+        u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space_a, name="tlm_u").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
         v = Function(space_b, name="v").project(u)
         J = assemble(v * v * dx)
 
@@ -161,11 +161,11 @@ def test_supermesh_project():
     space_b = FunctionSpace(mesh_b, "Discontinuous Lagrange", 0)
     test_a = TestFunction(space_a)
 
-    u = Function(space_a, name="u").interpolate(X_a[0] - 0.5)
-    zeta = Function(space_a, name="tlm_u").interpolate(X_a[0])
-    u.block_variable.tlm_value = zeta.copy(deepcopy=True)
-
     with reverse_over_forward():
+        u = Function(space_a, name="u").interpolate(X_a[0] - 0.5)
+        zeta = Function(space_a, name="tlm_u").interpolate(X_a[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
         v = Function(space_b, name="v").project(u)
         J = assemble(v * v * dx)
 

From e0666e5407bc4d8d1f00eed492ffb01d79722102 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 19:55:18 +0100
Subject: [PATCH 14/29] Reverse-over-forward AD: DirichletBCBlock

---
 .../adjoint_utils/blocks/dirichlet_bc.py      | 12 +++++++++++
 firedrake/adjoint_utils/blocks/solving.py     | 10 +++++-----
 .../test_adjoint_reverse_over_forward.py      | 20 +++++++++++++++++++
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
index 7a4784812c..aed39f1a35 100644
--- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py
+++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
@@ -121,6 +121,18 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
             m = bc.reconstruct(g=tlm_input)
         return m
 
+    def solve_tlm(self):
+        x, = self.get_outputs()
+        dep = x.output._function_arg
+        if isinstance(dep, OverloadedType):
+            tlm_dep = x.output._function_arg.block_variable.tlm_value
+        else:
+            tlm_dep = None
+        if tlm_dep is None:
+            x.tlm_value = None
+        else:
+            x.tlm_value = x.output.reconstruct(g=tlm_dep)
+
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx, relevant_dependencies,
                                    prepared=None):
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 2ca1c86c92..7158b3f7b2 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -345,12 +345,12 @@ def solve_tlm(self):
                     tlm_rhs = tlm_rhs - firedrake.action(
                         firedrake.derivative(form, dep), tlm_dep)
 
-        x.tlm_value = None
         if isinstance(tlm_rhs, int) and tlm_rhs == 0:
-            return
-        tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
-        if tlm_rhs.empty():
-            return
+            tlm_rhs = firedrake.Cofunction(x.output.function_space().dual())
+        else:
+            tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
+            if tlm_rhs.empty():
+                tlm_rhs = firedrake.Cofunction(x.output.function_space().dual())
 
         if self.linear:
             J = self.lhs
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 2cada1e4a2..9a368a9ba9 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -173,3 +173,23 @@ def test_supermesh_project():
     adj_value = u.block_variable.adj_value
     assert np.allclose(adj_value.dat.data_ro,
                        assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)
+
+
+
+@pytest.mark.skipcomplex
+def test_dirichletbc():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+
+    with reverse_over_forward():
+        u = Function(space, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+        bc = DirichletBC(space, u, "on_boundary")
+        
+        v = project(Constant(0.0), space, bcs=bc)
+        J = assemble(v * v * v * dx)
+
+    J_hat = ReducedFunctional(J.block_variable.tlm_value, Control(u))
+    assert taylor_test(J_hat, u, Function(space).interpolate(X[0] * X[0])) > 1.9

From c20303a0de4480bd3b3ecbfe9d808dcb02032beb Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 19:56:03 +0100
Subject: [PATCH 15/29] Fix names

---
 .../regression/test_adjoint_reverse_over_forward.py  | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 9a368a9ba9..55615b7272 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -36,7 +36,7 @@ def test_assembly():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0])
-        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         J = assemble(u * u * dx)
@@ -72,7 +72,7 @@ def test_function_assignment():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
-        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space, name="v").assign(u)
@@ -93,7 +93,7 @@ def test_function_assignment_expr():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
-        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space, name="v").assign(-3 * u)
@@ -140,7 +140,7 @@ def test_project():
 
     with reverse_over_forward():
         u = Function(space_a, name="u").interpolate(X[0] - 0.5)
-        zeta = Function(space_a, name="tlm_u").interpolate(X[0])
+        zeta = Function(space_a, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space_b, name="v").project(u)
@@ -163,7 +163,7 @@ def test_supermesh_project():
 
     with reverse_over_forward():
         u = Function(space_a, name="u").interpolate(X_a[0] - 0.5)
-        zeta = Function(space_a, name="tlm_u").interpolate(X_a[0])
+        zeta = Function(space_a, name="zeta").interpolate(X_a[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space_b, name="v").project(u)
@@ -184,7 +184,7 @@ def test_dirichletbc():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
-        zeta = Function(space, name="tlm_u").interpolate(X[0])
+        zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
         bc = DirichletBC(space, u, "on_boundary")
         

From 8d034026e898425418bd1b3e70ef52e1c78e55b3 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Wed, 10 Jul 2024 20:11:15 +0100
Subject: [PATCH 16/29] flake8

---
 tests/regression/test_adjoint_reverse_over_forward.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 55615b7272..7b95150af3 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -175,7 +175,6 @@ def test_supermesh_project():
                        assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)
 
 
-
 @pytest.mark.skipcomplex
 def test_dirichletbc():
     mesh = UnitIntervalMesh(10)
@@ -187,7 +186,7 @@ def test_dirichletbc():
         zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
         bc = DirichletBC(space, u, "on_boundary")
-        
+
         v = project(Constant(0.0), space, bcs=bc)
         J = assemble(v * v * v * dx)
 

From f48062b4f5e84af8acecb5b8424d7c19947cf251 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Thu, 11 Jul 2024 13:25:44 +0100
Subject: [PATCH 17/29] Test case where inputs and outputs are different
 variable versions

---
 firedrake/adjoint_utils/blocks/solving.py     | 10 +++++--
 .../test_adjoint_reverse_over_forward.py      | 28 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 7158b3f7b2..5ef8c4a9c6 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -320,15 +320,18 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
     def solve_tlm(self):
         x, = self.get_outputs()
         if self.linear:
-            form = firedrake.action(self.lhs, x.output) - self.rhs
+            tmp_x = firedrake.Function(x.output.function_space())
+            replace_map = {tmp_x: x.output}
+            form = firedrake.action(self.lhs, tmp_x) - self.rhs
         else:
+            replace_map = None
             form = self.lhs
 
         tlm_rhs = 0
         tlm_bcs = []
         for block_variable in self.get_dependencies():
             dep = block_variable.output
-            if dep == x.output:
+            if dep == x.output and not self.linear:
                 continue
             tlm_dep = block_variable.tlm_value
             if isinstance(dep, firedrake.DirichletBC):
@@ -357,6 +360,9 @@ def solve_tlm(self):
         else:
             J = firedrake.derivative(form, x.output, firedrake.TrialFunction(x.output.function_space()))
 
+        if replace_map is not None:
+            J = ufl.replace(J, replace_map)
+            tlm_rhs = ufl.replace(tlm_rhs, replace_map)
         x.tlm_value = firedrake.Function(x.output.function_space())
         firedrake.solve(J == tlm_rhs, x.tlm_value, tlm_bcs, *self.forward_args,
                         **self.forward_kwargs)
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 7b95150af3..3ecd03aa96 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -152,6 +152,34 @@ def test_project():
                        assemble(2 * inner(Function(space_b).project(zeta), test_a) * dx).dat.data_ro)
 
 
+@pytest.mark.skipcomplex
+@pytest.mark.xfail  # Firedrake issue #3682
+def test_project_overwrite():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space = FunctionSpace(mesh, "Lagrange", 1)
+    test = TestFunction(space)
+
+    with reverse_over_forward():
+        u = Function(space, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space, name="zeta").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+        v = Function(space).project(-2 * u)
+        u.project(-2 * u)
+        assert np.allclose(u.dat.data_ro, v.dat.data_ro)
+        assert np.allclose(v.block_variable.tlm_value.dat.data_ro,
+                           -2 * zeta.dat.data_ro)
+        assert np.allclose(u.block_variable.tlm_value.dat.data_ro,
+                           -2 * zeta.dat.data_ro)
+        J = assemble(u * u * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(8 * inner(zeta, test) * dx).dat.data_ro)
+
+
 @pytest.mark.skipcomplex
 def test_supermesh_project():
     mesh_a = UnitSquareMesh(10, 10)

From 386f1f90ae9c2c157acb1d5ee2cd2f5120eeb5cd Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Thu, 11 Jul 2024 13:55:46 +0100
Subject: [PATCH 18/29] Fix test_project_overwrite

---
 tests/regression/test_adjoint_reverse_over_forward.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 3ecd03aa96..fe6d737197 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -153,7 +153,6 @@ def test_project():
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.xfail  # Firedrake issue #3682
 def test_project_overwrite():
     mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
@@ -166,13 +165,14 @@ def test_project_overwrite():
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space).project(-2 * u)
-        u.project(-2 * u)
-        assert np.allclose(u.dat.data_ro, v.dat.data_ro)
+        w = Function(space).assign(u)
+        w.project(-2 * w)
+        assert np.allclose(w.dat.data_ro, v.dat.data_ro)
         assert np.allclose(v.block_variable.tlm_value.dat.data_ro,
                            -2 * zeta.dat.data_ro)
-        assert np.allclose(u.block_variable.tlm_value.dat.data_ro,
+        assert np.allclose(w.block_variable.tlm_value.dat.data_ro,
                            -2 * zeta.dat.data_ro)
-        J = assemble(u * u * dx)
+        J = assemble(w * w * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value

From 468de377ca1741f180c17fea2859dd00754fcdf8 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Thu, 11 Jul 2024 19:00:25 +0100
Subject: [PATCH 19/29] Import fix

---
 firedrake/adjoint_utils/blocks/constant.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/constant.py b/firedrake/adjoint_utils/blocks/constant.py
index 7b38a063fd..d384e2dcef 100644
--- a/firedrake/adjoint_utils/blocks/constant.py
+++ b/firedrake/adjoint_utils/blocks/constant.py
@@ -1,7 +1,8 @@
-import firedrake
-import numpy
 from pyadjoint import Block, OverloadedType
+import numpy
+
 from pyadjoint.reduced_functional_numpy import gather
+import firedrake
 from .block_utils import isconstant
 
 

From f37e6421a7c4a393f777091e554d41b7ae0a5fc9 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 12:59:42 +0100
Subject: [PATCH 20/29] Add TLM Taylor verification to test_dirichletbc

---
 tests/regression/test_adjoint_reverse_over_forward.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index fe6d737197..c66f496e56 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -218,5 +218,7 @@ def test_dirichletbc():
         v = project(Constant(0.0), space, bcs=bc)
         J = assemble(v * v * v * dx)
 
+    J_hat = ReducedFunctional(J, Control(u))
+    assert taylor_test(J_hat, u, zeta, dJdm=J.block_variable.tlm_value) > 1.9
     J_hat = ReducedFunctional(J.block_variable.tlm_value, Control(u))
     assert taylor_test(J_hat, u, Function(space).interpolate(X[0] * X[0])) > 1.9

From f5023c6875eeb4cc83fba1b10909fcec59c4d76f Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 13:14:20 +0100
Subject: [PATCH 21/29] Reverse-over-forward: AssembleBlock bugfix

---
 firedrake/adjoint_utils/blocks/assembly.py    |  2 +-
 .../test_adjoint_reverse_over_forward.py      | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/assembly.py b/firedrake/adjoint_utils/blocks/assembly.py
index 05b7672300..486e86ffaa 100644
--- a/firedrake/adjoint_utils/blocks/assembly.py
+++ b/firedrake/adjoint_utils/blocks/assembly.py
@@ -166,7 +166,7 @@ def solve_tlm(self):
         if isinstance(tlm_rhs, int) and tlm_rhs == 0:
             return
         tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
-        if tlm_rhs.empty():
+        if isinstance(tlm_rhs, ufl.ZeroBaseForm) or (isinstance(tlm_rhs, ufl.Form) and tlm_rhs.empty()):
             return
         x.tlm_value = firedrake.assemble(tlm_rhs)
 
diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index c66f496e56..7307aa3d83 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -130,6 +130,28 @@ def test_subfunction(idx):
                        assemble(2 * inner(zeta[idx], test[idx]) * dx).dat.data_ro[idx])
 
 
+@pytest.mark.skipcomplex
+def test_interpolate():
+    mesh = UnitSquareMesh(10, 10)
+    X = SpatialCoordinate(mesh)
+    space_a = FunctionSpace(mesh, "Lagrange", 1)
+    space_b = FunctionSpace(mesh, "Lagrange", 2)
+    test_a = TestFunction(space_a)
+
+    with reverse_over_forward():
+        u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+        zeta = Function(space_a, name="zeta").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+        v = Function(space_b, name="v").interpolate(u)
+        J = assemble(v * v * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(adj_value.dat.data_ro,
+                       assemble(2 * inner(zeta, test_a) * dx).dat.data_ro)
+
+
 @pytest.mark.skipcomplex
 def test_project():
     mesh = UnitSquareMesh(10, 10)

From bf343ab81ec9d91f1f8acadfd771ae7ad7b33583 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 13:19:01 +0100
Subject: [PATCH 22/29] Defensive copy to handle assemble(Cofunction) case

---
 firedrake/adjoint_utils/blocks/assembly.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/assembly.py b/firedrake/adjoint_utils/blocks/assembly.py
index 486e86ffaa..82568b6848 100644
--- a/firedrake/adjoint_utils/blocks/assembly.py
+++ b/firedrake/adjoint_utils/blocks/assembly.py
@@ -168,7 +168,10 @@ def solve_tlm(self):
         tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
         if isinstance(tlm_rhs, ufl.ZeroBaseForm) or (isinstance(tlm_rhs, ufl.Form) and tlm_rhs.empty()):
             return
-        x.tlm_value = firedrake.assemble(tlm_rhs)
+        tlm_rhs = firedrake.assemble(tlm_rhs)
+        if tlm_rhs in {dep.output for dep in self.get_dependencies()}:
+            tlm_rhs = tlm_rhs.copy(deepcopy=True)
+        x.tlm_value = tlm_rhs
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):

From d494d0c3c4fd60a0c7522b8003e9cc8bf672aa83 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 13:24:15 +0100
Subject: [PATCH 23/29] ConstantAssignBlock.solve_tlm tidying

---
 firedrake/adjoint_utils/blocks/constant.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/constant.py b/firedrake/adjoint_utils/blocks/constant.py
index d384e2dcef..40178916b7 100644
--- a/firedrake/adjoint_utils/blocks/constant.py
+++ b/firedrake/adjoint_utils/blocks/constant.py
@@ -72,19 +72,21 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
         return constant_from_values(block_variable.output, values)
 
     def solve_tlm(self):
-        x, = self.get_outputs()
-        if len(x.output.ufl_shape) == 0:
-            x.tlm_value = firedrake.Constant(0.0)
-        else:
-            x.tlm_value = firedrake.Constant(
-                numpy.reshape(numpy.zeros_like(x.output.values()), x.output.ufl_shape))
         if self.assigned_list:
             # Not reachable?
             raise NotImplementedError
+
+        x, = self.get_outputs()
+        dep, = self.get_dependencies()
+        if dep.tlm_value is None:
+            x.tlm_value = None
         else:
-            dep, = self.get_dependencies()
-            if dep.tlm_value is not None:
-                x.tlm_value.assign(dep.tlm_value)
+            if len(x.output.ufl_shape) == 0:
+                x.tlm_value = firedrake.Constant(0.0)
+            else:
+                x.tlm_value = firedrake.Constant(
+                    numpy.reshape(numpy.zeros_like(x.output.values()), x.output.ufl_shape))
+            x.tlm_value.assign(dep.tlm_value)
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
                                  relevant_dependencies):

From 0dc5269b262014f9a1e1c338f9ca46b4e42855cb Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 13:26:43 +0100
Subject: [PATCH 24/29] DirichletBCBlock.solve_tlm tidying

---
 firedrake/adjoint_utils/blocks/dirichlet_bc.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
index aed39f1a35..9ee708c5b6 100644
--- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py
+++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py
@@ -125,13 +125,9 @@ def solve_tlm(self):
         x, = self.get_outputs()
         dep = x.output._function_arg
         if isinstance(dep, OverloadedType):
-            tlm_dep = x.output._function_arg.block_variable.tlm_value
+            x.tlm_value = x.output.reconstruct(g=dep.block_variable.tlm_value)
         else:
-            tlm_dep = None
-        if tlm_dep is None:
             x.tlm_value = None
-        else:
-            x.tlm_value = x.output.reconstruct(g=tlm_dep)
 
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
                                    block_variable, idx, relevant_dependencies,

From 00ddc36088d5e833ee24501d6529baa245ae4d9e Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 13:34:49 +0100
Subject: [PATCH 25/29] Handle ZeroBaseForm case

---
 firedrake/adjoint_utils/blocks/solving.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 5ef8c4a9c6..72c3056dae 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -352,7 +352,7 @@ def solve_tlm(self):
             tlm_rhs = firedrake.Cofunction(x.output.function_space().dual())
         else:
             tlm_rhs = ufl.algorithms.expand_derivatives(tlm_rhs)
-            if tlm_rhs.empty():
+            if isinstance(tlm_rhs, ufl.ZeroBaseForm) or (isinstance(tlm_rhs, ufl.Form) and tlm_rhs.empty()):
                 tlm_rhs = firedrake.Cofunction(x.output.function_space().dual())
 
         if self.linear:

From 36883436eb11068f56371e30d555d5c132394ec2 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 14:01:02 +0100
Subject: [PATCH 26/29] Test improvements

---
 .../test_adjoint_reverse_over_forward.py      | 67 ++++++++++++-------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 7307aa3d83..dc28bef5f4 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -35,16 +35,18 @@ def test_assembly():
     test = TestFunction(space)
 
     with reverse_over_forward():
-        u = Function(space, name="u").interpolate(X[0])
+        u = Function(space, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
-        J = assemble(u * u * dx)
+        J = assemble((u ** 3) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(6 * inner(u_ref * zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -72,16 +74,18 @@ def test_function_assignment():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space, name="v").assign(u)
-        J = assemble(v * v * dx)
+        J = assemble((v ** 3) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(zeta, test) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(6 * inner(u_ref * zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -94,15 +98,17 @@ def test_function_assignment_expr():
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
         zeta = Function(space, name="zeta").interpolate(X[0])
+        u_ref = u.copy(deepcopy=True)
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space, name="v").assign(-3 * u)
-        J = assemble(v * v * dx)
+        J = assemble((v ** 3) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(18 * inner(zeta, test) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(-162 * inner(u_ref * zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -116,18 +122,20 @@ def test_subfunction(idx):
     with reverse_over_forward():
         u = Function(space, name="u")
         u.sub(idx).interpolate(-2 * X[0])
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space, name="zeta")
         zeta.sub(idx).interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space, name="v")
         v.sub(idx).assign(u.sub(idx))
-        J = assemble(u.sub(idx) * v.sub(idx) * dx)
+        J = assemble((u.sub(idx) ** 2) * v.sub(idx) * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.sub(idx).dat.data_ro,
-                       assemble(2 * inner(zeta[idx], test[idx]) * dx).dat.data_ro[idx])
+    assert np.allclose(
+        adj_value.sub(idx).dat.data_ro,
+        assemble(6 * inner(u_ref[idx] * zeta[idx], test[idx]) * dx).dat.data_ro[idx])
 
 
 @pytest.mark.skipcomplex
@@ -140,16 +148,18 @@ def test_interpolate():
 
     with reverse_over_forward():
         u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space_a, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space_b, name="v").interpolate(u)
-        J = assemble(v * v * dx)
+        J = assemble(v ** 3 * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(zeta, test_a) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(6 * inner(u_ref * zeta, test_a) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -162,16 +172,18 @@ def test_project():
 
     with reverse_over_forward():
         u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space_a, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space_b, name="v").project(u)
-        J = assemble(v * v * dx)
+        J = assemble(v ** 3 * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(Function(space_b).project(zeta), test_a) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(6 * inner(Function(space_b).project(u_ref) * Function(space_b).project(zeta), test_a) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -183,6 +195,7 @@ def test_project_overwrite():
 
     with reverse_over_forward():
         u = Function(space, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
         zeta = Function(space, name="zeta").interpolate(X[0])
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
@@ -194,12 +207,13 @@ def test_project_overwrite():
                            -2 * zeta.dat.data_ro)
         assert np.allclose(w.block_variable.tlm_value.dat.data_ro,
                            -2 * zeta.dat.data_ro)
-        J = assemble(w * w * dx)
+        J = assemble(w ** 3 * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(8 * inner(zeta, test) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(-48 * inner(u_ref * zeta, test) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -217,12 +231,13 @@ def test_supermesh_project():
         u.block_variable.tlm_value = zeta.copy(deepcopy=True)
 
         v = Function(space_b, name="v").project(u)
-        J = assemble(v * v * dx)
+        J = assemble(v ** 2 * dx)
 
     _ = compute_gradient(J.block_variable.tlm_value, Control(u))
     adj_value = u.block_variable.adj_value
-    assert np.allclose(adj_value.dat.data_ro,
-                       assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(2 * inner(Function(space_a).project(Function(space_b).project(zeta)), test_a) * dx).dat.data_ro)
 
 
 @pytest.mark.skipcomplex
@@ -238,7 +253,7 @@ def test_dirichletbc():
         bc = DirichletBC(space, u, "on_boundary")
 
         v = project(Constant(0.0), space, bcs=bc)
-        J = assemble(v * v * v * dx)
+        J = assemble(v ** 3 * dx)
 
     J_hat = ReducedFunctional(J, Control(u))
     assert taylor_test(J_hat, u, zeta, dJdm=J.block_variable.tlm_value) > 1.9

From 5dec19adaf373665e8b17c2d3a40c29184107ff4 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 14:16:37 +0100
Subject: [PATCH 27/29] Test improvements

---
 .../test_adjoint_reverse_over_forward.py      | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index dc28bef5f4..89ac7b8384 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -121,7 +121,7 @@ def test_subfunction(idx):
 
     with reverse_over_forward():
         u = Function(space, name="u")
-        u.sub(idx).interpolate(-2 * X[0])
+        u.sub(idx).interpolate(X[0] - 0.5)
         u_ref = u.copy(deepcopy=True)
         zeta = Function(space, name="zeta")
         zeta.sub(idx).interpolate(X[0])
@@ -140,7 +140,7 @@ def test_subfunction(idx):
 
 @pytest.mark.skipcomplex
 def test_interpolate():
-    mesh = UnitSquareMesh(10, 10)
+    mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
     space_a = FunctionSpace(mesh, "Lagrange", 1)
     space_b = FunctionSpace(mesh, "Lagrange", 2)
@@ -162,9 +162,33 @@ def test_interpolate():
         assemble(6 * inner(u_ref * zeta, test_a) * dx).dat.data_ro)
 
 
+@pytest.mark.skipcomplex
+def test_interpolate_expr():
+    mesh = UnitIntervalMesh(10)
+    X = SpatialCoordinate(mesh)
+    space_a = FunctionSpace(mesh, "Lagrange", 1)
+    space_b = FunctionSpace(mesh, "Lagrange", 2)
+    test_a = TestFunction(space_a)
+
+    with reverse_over_forward():
+        u = Function(space_a, name="u").interpolate(X[0] - 0.5)
+        u_ref = u.copy(deepcopy=True)
+        zeta = Function(space_a, name="zeta").interpolate(X[0])
+        u.block_variable.tlm_value = zeta.copy(deepcopy=True)
+
+        v = Function(space_b, name="v").interpolate(-3 * u)
+        J = assemble(v ** 3 * dx)
+
+    _ = compute_gradient(J.block_variable.tlm_value, Control(u))
+    adj_value = u.block_variable.adj_value
+    assert np.allclose(
+        adj_value.dat.data_ro,
+        assemble(-162 * inner(u_ref * zeta, test_a) * dx).dat.data_ro)
+
+
 @pytest.mark.skipcomplex
 def test_project():
-    mesh = UnitSquareMesh(10, 10)
+    mesh = UnitIntervalMesh(10)
     X = SpatialCoordinate(mesh)
     space_a = FunctionSpace(mesh, "Lagrange", 1)
     space_b = FunctionSpace(mesh, "Discontinuous Lagrange", 0)

From ab80fe6528456bfc3759b00b69823dd6c48ab01d Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Thu, 11 Jul 2024 18:37:21 +0100
Subject: [PATCH 28/29] DROP BEFORE MERGE

---
 .github/workflows/build.yml | 1 +
 requirements-git.txt        | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 4f61ab2d9b..f3c791e05f 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -82,6 +82,7 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
+            --package-branch pyadjoint jrmaddison/reverse_over_forward \
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |
diff --git a/requirements-git.txt b/requirements-git.txt
index 8bf05aad21..db584f08d8 100644
--- a/requirements-git.txt
+++ b/requirements-git.txt
@@ -3,6 +3,6 @@ git+https://github.com/firedrakeproject/fiat.git#egg=fiat
 git+https://github.com/FInAT/FInAT.git#egg=finat
 git+https://github.com/firedrakeproject/tsfc.git#egg=tsfc
 git+https://github.com/OP2/PyOP2.git#egg=pyop2
-git+https://github.com/dolfin-adjoint/pyadjoint.git#egg=pyadjoint
+git+https://github.com/jrmaddison/pyadjoint.git#egg=pyadjoint
 git+https://github.com/firedrakeproject/petsc.git@firedrake#egg=petsc
 git+https://github.com/firedrakeproject/pytest-mpi.git@main#egg=pytest-mpi

From 4bb4e3b52fd9bd712842e78e64fc86ef1cce0363 Mon Sep 17 00:00:00 2001
From: "James R. Maddison" <j.r.maddison@ed.ac.uk>
Date: Fri, 12 Jul 2024 15:02:45 +0100
Subject: [PATCH 29/29] Fix context manager

---
 tests/regression/test_adjoint_reverse_over_forward.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/regression/test_adjoint_reverse_over_forward.py b/tests/regression/test_adjoint_reverse_over_forward.py
index 89ac7b8384..00c7d4530e 100644
--- a/tests/regression/test_adjoint_reverse_over_forward.py
+++ b/tests/regression/test_adjoint_reverse_over_forward.py
@@ -22,9 +22,11 @@ def _():
 def reverse_over_forward():
     continue_annotation()
     continue_reverse_over_forward()
-    yield
-    pause_annotation()
-    pause_reverse_over_forward()
+    try:
+        yield
+    finally:
+        pause_annotation()
+        pause_reverse_over_forward()
 
 
 @pytest.mark.skipcomplex