From 48dc5454b9ca8e34e8d3ca47a57b5b5bebc4aec0 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Mon, 8 Jun 2026 11:42:13 +0100
Subject: [PATCH 1/8] Add CachedSolverBlock

---
 firedrake/adjoint_utils/blocks/__init__.py    |   4 +-
 firedrake/adjoint_utils/blocks/solving.py     | 390 ++++++++++++++++--
 firedrake/adjoint_utils/variational_solver.py | 367 +++++++++++++++-
 tests/firedrake/adjoint/test_assemble.py      |   8 +-
 tests/firedrake/adjoint/test_hessian.py       |  46 +--
 tests/firedrake/adjoint/test_nlvs.py          | 191 +++++++++
 6 files changed, 924 insertions(+), 82 deletions(-)
 create mode 100644 tests/firedrake/adjoint/test_nlvs.py

diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py
index 3bbc926ce9..9e4c294b39 100644
--- a/firedrake/adjoint_utils/blocks/__init__.py
+++ b/firedrake/adjoint_utils/blocks/__init__.py
@@ -1,7 +1,7 @@
 from firedrake.adjoint_utils.blocks.assembly import AssembleBlock  # noqa F401
 from firedrake.adjoint_utils.blocks.solving import (  # noqa F401
-    GenericSolveBlock, SolveLinearSystemBlock, ProjectBlock,
-    SupermeshProjectBlock, SolveVarFormBlock,
+    CachedSolverBlock, GenericSolveBlock, SolveLinearSystemBlock,
+    ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock,
     NonlinearVariationalSolveBlock
 )
 from firedrake.adjoint_utils.blocks.function import (  # noqa F401
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 1d46f8a488..2c75fe9cb3 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -25,10 +25,296 @@ def extract_subfunction(u, V):
         return u
 
 
-class Solver(Enum):
+class SolverType(Enum):
     """Enum for solver types."""
     FORWARD = 0
     ADJOINT = 1
+    TLM = 2
+    HESSIAN = 3
+
+
+FORWARD = SolverType.FORWARD
+ADJOINT = SolverType.ADJOINT
+TLM = SolverType.TLM
+HESSIAN = SolverType.HESSIAN
+
+
+class CachedSolverBlock(Block):
+    def __init__(self, func, bcs, cached_solvers,
+                 is_linear, replaced_dependencies,
+                 tlm_rhs, replaced_tlms, tlm_dFdm_forms,
+                 adj_rhs, adj_dFdm_forms, adj_residual,
+                 adj_sol, adj2_sol, tlm_output,
+                 d2Fdu2_form, d2Fdmdu_forms,
+                 dFdm_adj2_forms, d2Fdm2_adj_forms,
+                 d2Fdudm_forms,
+                 ad_block_tag=None):
+        super().__init__(ad_block_tag=ad_block_tag)
+
+        self.func = func
+        self.bcs = bcs
+        self.cached_solvers = cached_solvers
+        self.replaced_dependencies = replaced_dependencies
+        self.is_linear = is_linear
+
+        self.tlm_rhs = tlm_rhs
+        self.replaced_tlms = replaced_tlms
+        self.tlm_dFdm_forms = tlm_dFdm_forms
+
+        self.adj_rhs = adj_rhs
+        self.adj_dFdm_forms = adj_dFdm_forms
+        self.adj_residual = adj_residual
+
+        self.adj_sol = adj_sol
+        self.adj_sol_buf = adj_sol.copy(deepcopy=True)
+        self.adj2_sol = adj2_sol
+        self.tlm_output = tlm_output
+        self.d2Fdu2_form = d2Fdu2_form
+        self.d2Fdmdu_forms = d2Fdmdu_forms
+        self.dFdm_adj2_forms = dFdm_adj2_forms
+        self.d2Fdm2_adj_forms = d2Fdm2_adj_forms
+        self.d2Fdudm_forms = d2Fdudm_forms
+
+    def _coefficient_dependencies(self, dependencies=None):
+        dependencies = dependencies or self.get_dependencies()
+        return dependencies[:len(self.replaced_dependencies)]
+
+    def _bc_dependencies(self, dependencies=None):
+        dependencies = dependencies or self.get_dependencies()
+        if len(self.bcs) > 0:
+            return dependencies[-len(self.bcs):]
+        else:
+            return []
+
+    def update_dependencies(self, use_output=False):
+        """Update all dependencies of the forward solve.
+        """
+        # Update the coefficients in the form.
+        # Use the fact that zip will use the shorter length.
+        for replaced_dep, dep in zip(self.replaced_dependencies,
+                                     self._coefficient_dependencies()):
+            replaced_dep.assign(dep.saved_output)
+
+        # 1. For forward recomputation the unknown Function should use
+        # the incoming value of the dependency as the initial guess.
+        # 2. For the adjoint, TLM, and Hessian, the unknown Function
+        # should use the computed value so that the linearised
+        # Jacobian is correct.
+        if use_output:
+            output = self.get_outputs()[0].saved_output
+            self.cached_solvers[FORWARD]._problem.u.assign(output)
+
+        # Update the boundary conditions
+        for replaced_dep, dep in zip(self.bcs, self._bc_dependencies()):
+            replaced_dep.set_value(dep.saved_output.function_arg)
+
+    def update_tlm_dependencies(self):
+        """Update all dependencies of the tlm solve.
+        """
+        for replaced_dep, dep in zip(self.replaced_tlms,
+                                     self._coefficient_dependencies()):
+            if dep.output == self.func and not self.is_linear:
+                continue
+            if dep.tlm_value is None:  # This dependency doesn't depend on the controls
+                continue
+            replaced_dep.assign(dep.tlm_value)
+
+        for replaced_dep, dep in zip(self.bcs, self._bc_dependencies()):
+            if dep.tlm_value is None:  # This dependency doesn't depend on the controls
+                bc_val = 0
+            else:
+                bc_val = dep.tlm_value.function_arg
+            replaced_dep.set_value(bc_val)
+
+    def update_adj_dependencies(self):
+        # TODO: Anything to do here?
+        pass
+
+    def update_hessian_dependencies(self):
+        # TODO: Anything else to do here?
+        self.update_tlm_dependencies()
+
+    def _compute_boundary(self, relevant_dependencies):
+        return any(isinstance(dep.output, firedrake.DirichletBC)
+                   for _, dep in relevant_dependencies)
+
+    def prepare_recompute_component(self, inputs, relevant_outputs):
+        return
+
+    def recompute_component(self, inputs, block_variable, idx, prepared):
+        self.update_dependencies(use_output=False)
+
+        solver = self.cached_solvers[FORWARD]
+        solver.solve()
+        result = solver._problem.u.copy(deepcopy=True)
+
+        # Possibly checkpoint the result for the adjoint solve later.
+        if isinstance(block_variable.checkpoint, firedrake.Function):
+            result = block_variable.checkpoint.assign(result)
+
+        return maybe_disk_checkpoint(result)
+
+    def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
+        return
+
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
+        self.update_dependencies(use_output=True)
+        self.update_tlm_dependencies()
+
+        # Assemble the rhs of (dF/du)(du/dm) = -dF/dm
+        self.tlm_rhs.zero()
+        for dFdm, dep in zip(self.tlm_dFdm_forms, self.get_dependencies()):
+            if dep.tlm_value is None:  # This dependency doesn't depend on the controls
+                continue
+            if dep.output is self.func and not self.is_linear:  # Can't compute dependence on initial guess
+                continue
+            self.tlm_rhs += firedrake.assemble(dFdm)
+
+        # Solve for dudm
+        solver = self.cached_solvers[TLM]
+        solver._problem.u.zero()
+        solver.solve()
+        result = solver._problem.u.copy(deepcopy=True)
+        return result
+
+    def solve_adj_equation(self, rhs, compute_boundary):
+        for bc in self.bcs:
+            bc.homogenize()
+
+        solver = self.cached_solvers[ADJOINT]
+        adj_sol = solver._problem.u
+
+        self.adj_rhs.assign(rhs)
+        adj_sol.zero()
+
+        solver.solve()
+
+        self.adj_sol.assign(adj_sol)
+        self.adj_sol_buf.assign(adj_sol)
+
+        if compute_boundary:
+            adj_sol_bc = firedrake.assemble(self.adj_residual)
+            adj_sol_bc = adj_sol_bc.riesz_representation("l2")
+        else:
+            adj_sol_bc = None
+
+        return adj_sol, adj_sol_bc
+
+    def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
+        self.update_dependencies(use_output=True)
+        self.update_adj_dependencies()
+
+        dJdu = adj_inputs[0]
+
+        compute_boundary = self._compute_boundary(relevant_dependencies)
+
+        adj_sol, adj_sol_bc = self.solve_adj_equation(dJdu, compute_boundary)
+
+        # store adj_sol for Hessian computation later.
+        # self.adj_sol is shared between all blocks that this NLVS
+        # generates so we can't store it there. Instead store it
+        # in self.adj_sol_buf which is owned by this block only.
+        self.adj_sol_buf.assign(adj_sol)
+
+        prepared = {
+            "adj_sol": adj_sol.copy(deepcopy=True),
+            "adj_sol_bc": adj_sol_bc
+        }
+        return prepared
+
+    def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
+        if block_variable.output == self.func and not self.is_linear:
+            return None
+
+        if isinstance(block_variable.output, firedrake.DirichletBC):
+            bc = block_variable.output
+            adj_sol_bc = prepared["adj_sol_bc"]
+            return bc.reconstruct(
+                g=extract_subfunction(adj_sol_bc, bc.function_space())
+            )
+
+        # assemble sensititivy comment
+        dFdm = firedrake.assemble(self.adj_dFdm_forms[idx])
+
+        return dFdm
+
+    def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
+        self.adj_sol.assign(self.adj_sol_buf)
+        self.update_dependencies(use_output=True)
+        self.update_hessian_dependencies()
+
+        hessian_input = hessian_inputs[0]
+        tlm_output = self.get_outputs()[0].tlm_value
+
+        if hessian_input is None:
+            return
+        if tlm_output is None:
+            return
+
+        # 1. Assemble rhs
+
+        # hessian input contribution
+        hessian_rhs = hessian_input.copy(deepcopy=True)
+
+        # tlm_output contribution
+        self.tlm_output.assign(tlm_output)
+        hessian_rhs -= firedrake.assemble(self.d2Fdu2_form)
+
+        # tlm_input contribution
+        for d2Fdmdu, dep in zip(self.d2Fdmdu_forms,
+                                self._coefficient_dependencies()):
+            if dep.tlm_value is None:  # This dependency doesn't depend on the controls
+                continue
+            if dep.output is self.func and not self.is_linear:  # Can't compute dependence on initial guess
+                continue
+            if len(d2Fdmdu.integrals()) > 0:
+                hessian_rhs -= firedrake.assemble(d2Fdmdu)
+
+        # 2. Solve adjoint system
+        compute_boundary = self._compute_boundary(relevant_dependencies)
+        adj2_sol, adj2_sol_bc = self.solve_adj_equation(hessian_rhs, compute_boundary)
+
+        self.adj2_sol.assign(adj2_sol)
+
+        prepared = {
+            "adj2_sol": adj2_sol.copy(deepcopy=True),
+            "adj2_sol_bc": adj2_sol_bc,
+        }
+
+        return prepared
+
+    def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None):
+        m = block_variable.output
+
+        if m is self.func and not self.is_linear:
+            return None
+
+        if isinstance(m, firedrake.DirichletBC):
+            bc = block_variable.output
+            adj2_sol_bc = prepared["adj2_sol_bc"]
+            return bc.reconstruct(
+                g=extract_subfunction(adj2_sol_bc, bc.function_space())
+            )
+
+        relevant_d2Fdm2_forms = []
+        for i, dep in relevant_dependencies:
+            if i >= len(self._coefficient_dependencies()):
+                continue
+            if dep.tlm_value is None:
+                continue
+            if dep.output is self.func and not self.is_linear:
+                continue
+            relevant_d2Fdm2_forms.append(self.d2Fdm2_adj_forms[idx][i])
+
+        hessian_output = 0
+
+        for form in (self.d2Fdudm_forms[idx],
+                     self.dFdm_adj2_forms[idx],
+                     *relevant_d2Fdm2_forms):
+            if not form.empty():
+                hessian_output += firedrake.assemble(-form)
+
+        return hessian_output
 
 
 class GenericSolveBlock(Block):
@@ -56,6 +342,8 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         # Solution function
         self.func = func
         self.function_space = self.func.function_space()
+        # Storage for adjoint solution of this block
+        self.adj_state_buf = func.copy(deepcopy=True)
         # Boundary conditions
         self.bcs = []
         if bcs is not None:
@@ -204,6 +492,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             dFdu_form, dJdu, compute_bdy
         )
         self.adj_sol = adj_sol
+        self.adj_state = adj_sol
+        self.adj_state_buf.assign(adj_sol)
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
@@ -215,31 +505,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         r["adj_sol_bdy"] = adj_sol_bdy
         return r
 
-    def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
-        return firedrake.assemble(dFdu_adj_form, **kwargs)
-
-    def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
-        dJdu_copy = dJdu.copy()
-        # Homogenize and apply boundary conditions on adj_dFdu.
-        bcs = self._homogenize_bcs()
-        dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs)
-
-        adj_sol = firedrake.Function(self.function_space)
-        firedrake.solve(
-            dFdu, adj_sol, dJdu, *self.adj_args, **self.adj_kwargs
-        )
-
-        adj_sol_bdy = None
-        if compute_bdy:
-            adj_sol_bdy = self._compute_adj_bdy(
-                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
-
-        return adj_sol, adj_sol_bdy
-
-    def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
-        adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol))
-        return adj_sol_bdy.riesz_representation("l2")
-
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                                prepared=None):
         if not self.linear and self.func == block_variable.output:
@@ -283,7 +548,36 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs)
         return dFdm
 
+    def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
+        return firedrake.assemble(dFdu_adj_form, **kwargs)
+
+    def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
+        dJdu_copy = dJdu.copy()
+        # Homogenize and apply boundary conditions on adj_dFdu.
+        bcs = self._homogenize_bcs()
+        dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs)
+
+        adj_sol = firedrake.Function(self.function_space)
+        firedrake.solve(
+            dFdu, adj_sol, dJdu, *self.adj_args, **self.adj_kwargs
+        )
+
+        adj_sol_bdy = None
+        if compute_bdy:
+            adj_sol_bdy = self._compute_adj_bdy(
+                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
+
+        return adj_sol, adj_sol_bdy
+
+    def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
+        adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol))
+        return adj_sol_bdy.riesz_representation("l2")
+
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
+        pass
+
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
+                               prepared=None):
         fwd_block_variable = self.get_outputs()[0]
         u = fwd_block_variable.output
 
@@ -295,16 +589,6 @@ def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
             fwd_block_variable.saved_output,
             firedrake.TrialFunction(u.function_space())
         )
-
-        return {
-            "form": F_form,
-            "dFdu": dFdu
-        }
-
-    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
-                               prepared=None):
-        F_form = prepared["form"]
-        dFdu = prepared["dFdu"]
         V = self.get_outputs()[idx].output.function_space()
 
         bcs = []
@@ -341,10 +625,11 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx,
         dFdm = ufl.algorithms.expand_derivatives(dFdm)
         dFdm = firedrake.assemble(dFdm)
         dudm = firedrake.Function(V)
-        return self._assemble_and_solve_tlm_eq(
+        result = self._assemble_and_solve_tlm_eq(
             firedrake.assemble(dFdu, bcs=bcs, **self.assemble_kwargs),
             dFdm, dudm, bcs
         )
+        return result
 
     def _assemble_and_solve_tlm_eq(self, dFdu, dFdm, dudm, bcs):
         return self._assembled_solve(dFdu, dFdm, dudm, bcs)
@@ -376,7 +661,10 @@ def _assemble_soa_eq_rhs(self, dFdu_form, adj_sol, hessian_input, d2Fdu2):
             elif not isinstance(c, firedrake.DirichletBC):
                 dFdu_adj = firedrake.action(firedrake.adjoint(dFdu_form),
                                             adj_sol)
-                b_form += firedrake.derivative(dFdu_adj, c_rep, tlm_input)
+                # b_form += firedrake.derivative(dFdu_adj, c_rep, tlm_input)
+                bo_form = ufl.algorithms.expand_derivatives(
+                    firedrake.derivative(dFdu_adj, c_rep, tlm_input))
+                b_form += bo_form
 
         b_form = ufl.algorithms.expand_derivatives(b_form)
         if len(b_form.integrals()) > 0:
@@ -404,6 +692,8 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
         hessian_input = hessian_inputs[0]
         tlm_output = fwd_block_variable.tlm_value
 
+        self.adj_state = self.adj_state_buf.copy(deepcopy=True)
+
         if hessian_input is None:
             return
 
@@ -432,6 +722,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs,
         r["adj_sol2_bdy"] = adj_sol2_bdy
         r["form"] = F_form
         r["adj_sol"] = adj_sol
+
         return r
 
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
@@ -577,6 +868,12 @@ def solve_init_params(self, args, kwargs, varform):
                 )
             self.adj_kwargs.pop("appctx", None)
 
+    if hasattr(self, "tlm_args") and len(self.tlm_args) <= 0:
+        self.tlm_args = self.adj_args
+
+    if hasattr(self, "tlm_kwargs") and len(self.tlm_kwargs) <= 0:
+        self.tlm_kwargs = self.adj_kwargs.copy()
+
     solver_params = kwargs.get("solver_parameters", None)
     if solver_params is not None and "mat_type" in solver_params:
         self.assemble_kwargs["mat_type"] = solver_params["mat_type"]
@@ -627,6 +924,8 @@ def __init__(self, equation, func, bcs, adj_cache, problem_J,
         self.problem_J = problem_J
         self.solver_kwargs = solver_kwargs
 
+        self.adj_state_buf = func.copy(deepcopy=True)
+
         super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs})
 
         if self.problem_J is not None:
@@ -664,12 +963,12 @@ def _adjoint_solve(self, dJdu, compute_bdy):
             and self._ad_solvers["update_adjoint"]
         ):
             # Update left hand side of the adjoint equation.
-            self._ad_solver_replace_forms(Solver.ADJOINT)
+            self._ad_solver_replace_forms(SolverType.ADJOINT)
             self._ad_solvers["adjoint_lvs"].invalidate_jacobian()
             self._ad_solvers["update_adjoint"] = False
         elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
             # Update left hand side of the adjoint equation.
-            self._ad_solver_replace_forms(Solver.ADJOINT)
+            self._ad_solver_replace_forms(SolverType.ADJOINT)
 
         # Update the right hand side of the adjoint equation.
         # problem.F._component[1] is the right hand side of the adjoint.
@@ -687,7 +986,7 @@ def _adjoint_solve(self, dJdu, compute_bdy):
         return u_sol, adj_sol_bdy
 
     def _ad_assign_map(self, form, solver):
-        if solver == Solver.FORWARD:
+        if solver == SolverType.FORWARD:
             count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
         else:
             count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
@@ -705,7 +1004,7 @@ def _ad_assign_map(self, form, solver):
                         block_variable.saved_output
 
         if (
-            solver == Solver.ADJOINT
+            solver == SolverType.ADJOINT
             and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
         ):
             block_variable = self.get_outputs()[0]
@@ -720,8 +1019,8 @@ def _ad_assign_coefficients(self, form, solver):
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
-    def _ad_solver_replace_forms(self, solver=Solver.FORWARD):
-        if solver == Solver.FORWARD:
+    def _ad_solver_replace_forms(self, solver=SolverType.FORWARD):
+        if solver == SolverType.FORWARD:
             problem = self._ad_solvers["forward_nlvs"]._problem
             self._ad_assign_coefficients(problem.F, solver)
             self._ad_assign_coefficients(problem.J, solver)
@@ -734,7 +1033,8 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             relevant_dependencies
         )
         adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
-        self.adj_sol = adj_sol
+        self.adj_state = adj_sol
+        self.adj_state_buf.assign(adj_sol)
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 21dded9839..b5afc05733 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -1,9 +1,13 @@
 import copy
 from functools import wraps
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
-from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock
-from firedrake.ufl_expr import derivative, adjoint
-from ufl import replace
+from firedrake.adjoint_utils.blocks import (
+    NonlinearVariationalSolveBlock, CachedSolverBlock)
+from firedrake.adjoint_utils.blocks.solving import solve_init_params
+from firedrake.ufl_expr import derivative, adjoint, action
+from ufl import replace, Action
+from ufl.algorithms import expand_derivatives
+from types import SimpleNamespace
 
 
 class NonlinearVariationalProblemMixin:
@@ -52,10 +56,367 @@ def wrapper(self, problem, *args, **kwargs):
                                 "recompute_count": 0}
             self._ad_adj_cache = {}
 
+            self._ad_solver_cache = {}
+
+            # process args/kwargs for cached solvers
+            self._ad_args_kwargs = SimpleNamespace(
+                forward_args=kwargs.pop("forward_args", []),
+                forward_kwargs=kwargs.pop("forward_kwargs", {}),
+                adj_args=kwargs.pop("adj_args", []),
+                adj_kwargs=kwargs.pop("adj_kwargs", {}),
+                tlm_args=kwargs.pop("tlm_args", []),
+                tlm_kwargs=kwargs.pop("tlm_kwargs", {}),
+                assemble_kwargs={}
+            )
+            solve_init_params(self._ad_args_kwargs,
+                              args, kwargs, varform=True)
+
         return wrapper
 
+    def _ad_cache_forward_solver(self):
+        from firedrake import (
+            DirichletBC,
+            NonlinearVariationalProblem,
+            NonlinearVariationalSolver)
+        from firedrake.adjoint_utils.blocks.solving import FORWARD
+
+        problem = self._ad_problem
+
+        # Build a new form so that we can update the coefficient
+        # values without affecting user code.
+        # We do this by copying all coefficients in the form and
+        # symbolically replacing the old values with the new.
+        F = problem.F
+        replace_map = {}
+        for old_coeff in F.coefficients():
+            new_coeff = old_coeff.copy(deepcopy=True)
+            replace_map[old_coeff] = new_coeff
+
+        # We need a handle to the new Function being
+        # solved for so that we can create an NLVS.
+        Fnew = replace(F, replace_map)
+        unew = replace_map[problem.u]
+
+        # We also need to "replace" all the bcs in
+        # the new NLVS so we can modify those values
+        # without affecting user code.
+        # Note that ``DirichletBC.reconstruct`` will
+        # return ``self`` if V, g, and sub_domain are
+        # all unchanged, so we need to explicitly
+        # instantiate a new object.
+        bcs = problem.bcs
+        bcs_new = [
+            DirichletBC(V=bc.function_space(),
+                        g=bc.function_arg,
+                        sub_domain=bc.sub_domain)
+            for bc in bcs
+        ]
+
+        # This NLVS will be used to recompute the solve.
+        # TODO: solver_parameters
+        nlvp = NonlinearVariationalProblem(Fnew, unew, bcs=bcs_new)
+        nlvs = NonlinearVariationalSolver(
+            nlvp,
+            *self._ad_args_kwargs.forward_args,
+            **self._ad_args_kwargs.forward_kwargs)
+
+        # The original coefficients will be added as
+        # dependencies to all solve blocks.
+        # The block need handles to the newly created
+        # objects to update their values when recomputing.
+        self._ad_dependencies_to_add = (*replace_map.keys(), *bcs)
+        self._ad_replaced_dependencies = tuple(replace_map.values())
+        self._ad_bcs = bcs_new
+        self._ad_solver_cache[FORWARD] = nlvs
+
+    def _ad_cache_tlm_solver(self):
+        from firedrake import (
+            Function, Cofunction, derivative, TrialFunction,
+            LinearVariationalProblem, LinearVariationalSolver)
+        from firedrake.adjoint_utils.blocks.solving import FORWARD, TLM
+
+        # If we build the TLM form from the cached
+        # forward solve form then we can update exactly
+        # the same coefficients/boundary conditions.
+        nlvp = self._ad_solver_cache[FORWARD]._problem
+
+        F = nlvp.F
+        u = nlvp.u
+        V = u.function_space()
+
+        # We need gradient of output/input i.e. du/dm.
+        # We know F(u; m) = 0 and _total_ dF/dm = 0.
+        # Then for the _partial_ derivatives:
+        # (dF/du)*(du/dm) + dF/dm = 0 so we calculate:
+        # (dF/du)*(du/dm) = -dF/dm
+        dFdu = derivative(F, u, TrialFunction(V))
+        dFdm = Cofunction(V.dual())
+        dudm = Function(V)
+
+        self._ad_dFdu = dFdu
+
+        # Reuse the same bcs as the forward problem.
+        # TODO: Think about if we should use new bcs.
+        # TODO: solver_parameters
+        lvp = LinearVariationalProblem(dFdu, dFdm, dudm, bcs=self._ad_bcs)
+        lvs = LinearVariationalSolver(
+            lvp,
+            *self._ad_args_kwargs.tlm_args,
+            **self._ad_args_kwargs.tlm_kwargs)
+
+        self._ad_solver_cache[TLM] = lvs
+        self._ad_tlm_rhs = dFdm
+
+        # Do all the symbolic work for calculating dF/dm up front
+        # so we only pay for the numeric calculations at run time.
+        replaced_tlms = []
+        dFdm_tlm_forms = []
+        for m in self._ad_replaced_dependencies:
+            mtlm = m.copy(deepcopy=True)
+            replaced_tlms.append(mtlm)
+
+            dFdm = derivative(-F, m, mtlm)
+            # TODO: Do we need expand_derivatives here? If so, why?
+            dFdm = expand_derivatives(dFdm)
+            dFdm_tlm_forms.append(dFdm)
+
+        # We'll need to update the replaced_tlm
+        # values and assemble the dFdm forms
+        self._ad_tlm_dFdm_forms = dFdm_tlm_forms
+        self._ad_replaced_tlms = replaced_tlms
+
+    def _ad_cache_adj_solver(self):
+        from firedrake import (
+            Function, Cofunction, TrialFunction, Argument,
+            LinearVariationalProblem, LinearVariationalSolver)
+        from firedrake.adjoint_utils.blocks.solving import FORWARD, ADJOINT
+
+        # If we build the adjoint form from the cached
+        # forward solve form then we can update exactly
+        # the same coefficients/boundary conditions.
+        nlvp = self._ad_solver_cache[FORWARD]._problem
+
+        F = nlvp.F
+        u = nlvp.u
+        V = u.function_space()
+
+        # TODO: rewrite for adjoint not TLM
+        # We need gradient of output/input i.e. du/dm.
+        # We know F(u; m) = 0 and _total_ dF/dm = 0.
+        # Then for the _partial_ derivatives:
+        # (dF/du)*(du/dm) + dF/dm = 0 so we calculate:
+        # (dF/du)*(du/dm) = -dF/dm
+        dFdu = self._ad_dFdu
+        try:
+            dFdu_adj = adjoint(dFdu)
+        except ValueError:
+            # Try again without expanding derivatives,
+            # as dFdu might have been simplied to an empty Form
+            dFdu_adj = adjoint(dFdu, derivatives_expanded=True)
+
+        self._ad_dFdu_adj = dFdu_adj
+
+        # This will be the rhs of the adjoint problem
+        dJdu = Cofunction(V.dual())
+        adj_sol = Function(V)
+
+        # Reuse the same bcs as the forward problem.
+        # TODO: Think about if we should use new bcs.
+        # TODO: solver_parameters
+        lvp = LinearVariationalProblem(dFdu_adj, dJdu, adj_sol, bcs=self._ad_bcs)
+        lvs = LinearVariationalSolver(
+            lvp,
+            *self._ad_args_kwargs.adj_args,
+            **self._ad_args_kwargs.adj_kwargs)
+
+        self._ad_solver_cache[ADJOINT] = lvs
+        self._ad_adj_rhs = dJdu
+
+        # Do all the symbolic work for calculating dJ/du up front
+        # so we only pay for the numeric calculations at run time.
+        dFdm_adj_forms = []
+        for m in self._ad_replaced_dependencies:
+            # Action of adjoint solution on dFdm
+            # TODO: Which of the two implementations should we use?
+            dFdm = derivative(-F, m, TrialFunction(m.function_space()))
+
+            # 1. from previous cached implementation
+            dFdm = adjoint(dFdm)
+            if isinstance(dFdm, Argument):
+                #  Corner case. Should be fixed more permanently upstream in UFL.
+                #  See: https://github.com/FEniCS/ufl/issues/395
+                dFdm = Action(dFdm, adj_sol)
+            else:
+                dFdm = dFdm * adj_sol
+
+            # 2. from GenericSolveBlock
+            # if isinstance(dFdm, ufl.Form):
+            #     dFdm = adjoint(dFdm)
+            #     dFdm = action(dFdm, adj_sol)
+            # else:
+            #     dFdm = dFdm(adj_sol)
+
+            dFdm_adj_forms.append(dFdm)
+
+        # To calculate the adjoint component of each DirichletBC
+        # we'll need the residual of the adjoint equation without
+        # any DirichletBC using the solution calculated with
+        # homogeneous DirichletBCs.
+        self._ad_adj_residual = dJdu - action(dFdu_adj, adj_sol)
+
+        # We'll need to assemble these forms to calculate
+        # the adj_component for each dependency.
+        self._ad_adj_dFdm_forms = dFdm_adj_forms
+
+    def _ad_cache_hessian_solver(self):
+        from firedrake import (
+            Function, TestFunction)
+        from firedrake.adjoint_utils.blocks.solving import FORWARD
+
+        nlvp = self._ad_solver_cache[FORWARD]._problem
+        F = nlvp.F
+        u = nlvp.u
+        V = u.function_space()
+
+        # 1. Forms to calculate rhs of Hessian solve
+
+        # Calculate d^2F/du^2 * du/dm * dm
+        # where dm is direction for tlm action so du/dm * dm is tlm output
+        dFdu = self._ad_dFdu
+        tlm_output = Function(V)
+        d2Fdu2 = derivative(dFdu, u, tlm_output)
+        # print()
+        # print(f"{dFdu = }")
+        # print()
+        # print(f"{d2Fdu2 = }")
+        # print()
+        d2Fdu2 = expand_derivatives(d2Fdu2)
+
+        self._ad_tlm_output = tlm_output
+
+        adj_sol = Function(V)
+        self._ad_adj_sol = adj_sol
+
+        # Contribution from tlm_output
+        if len(d2Fdu2.integrals()) > 0:
+            d2Fdu2_form = action(adjoint(d2Fdu2), adj_sol)
+        else:
+            d2Fdu2_form = d2Fdu2
+        self._ad_d2Fdu2_form = d2Fdu2_form
+
+        # Contributions from each tlm_input
+        dFdu_adj = action(self._ad_dFdu_adj, adj_sol)
+        d2Fdmdu_forms = []
+        for m, dm in zip(self._ad_replaced_dependencies,
+                         self._ad_replaced_tlms):
+            d2Fdmdu = expand_derivatives(
+                derivative(dFdu_adj, m, dm))
+
+            d2Fdmdu_forms.append(d2Fdmdu)
+
+        self._ad_d2Fdmdu_forms = d2Fdmdu_forms
+
+        # 2. Forms to calculate contribution from each control
+        adj2_sol = Function(V)
+        self._ad_adj2_sol = adj2_sol
+
+        Fadj = action(F, adj_sol)
+        Fadj2 = action(F, adj2_sol)
+
+        dFdm_adj2_forms = []
+        d2Fdm2_adj_forms = []
+        d2Fdudm_forms = []
+        for m in self._ad_replaced_dependencies:
+            dm = TestFunction(m.function_space())
+            dFdm_adj2 = expand_derivatives(
+                derivative(Fadj2, m, dm))
+
+            dFdm_adj2_forms.append(dFdm_adj2)
+
+            dFdm_adj = derivative(Fadj, m, dm)
+
+            d2Fdudm = expand_derivatives(
+                derivative(dFdm_adj, u, tlm_output))
+
+            d2Fdudm_forms.append(d2Fdudm)
+
+            d2Fdm2_adj_forms_k = []
+            for m2, dm2 in zip(self._ad_replaced_dependencies,
+                               self._ad_replaced_tlms):
+                d2Fdm2_adj = expand_derivatives(
+                    derivative(dFdm_adj, m2, dm2))
+                d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
+
+            d2Fdm2_adj_forms.append(d2Fdm2_adj_forms_k)
+
+        self._ad_dFdm_adj2_forms = dFdm_adj2_forms
+        self._ad_d2Fdm2_adj_forms = d2Fdm2_adj_forms
+        self._ad_d2Fdudm_forms = d2Fdudm_forms
+
     @staticmethod
     def _ad_annotate_solve(solve):
+        @wraps(solve)
+        def wrapper(self, **kwargs):
+            """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
+            Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
+            for the purposes of the adjoint computation (such as projecting fields to other function spaces
+            for the purposes of visualisation)."""
+            annotate = annotate_tape(kwargs)
+            if annotate:
+                if kwargs.pop("bounds", None) is not None:
+                    raise ValueError(
+                        "MissingMathsError: we do not know how to differentiate through a variational inequality")
+
+                if len(self._ad_solver_cache) == 0:
+                    self._ad_cache_forward_solver()
+                    self._ad_cache_tlm_solver()
+                    self._ad_cache_adj_solver()
+                    self._ad_cache_hessian_solver()
+
+                block = CachedSolverBlock(self._ad_problem.u,
+                                          self._ad_bcs,
+                                          self._ad_solver_cache,
+                                          self._ad_problem.is_linear,
+                                          self._ad_replaced_dependencies,
+
+                                          self._ad_tlm_rhs,
+                                          self._ad_replaced_tlms,
+                                          self._ad_tlm_dFdm_forms,
+
+                                          self._ad_adj_rhs,
+                                          self._ad_adj_dFdm_forms,
+                                          self._ad_adj_residual,
+
+                                          self._ad_adj_sol,
+                                          self._ad_adj2_sol,
+                                          self._ad_tlm_output,
+                                          self._ad_d2Fdu2_form,
+                                          self._ad_d2Fdmdu_forms,
+                                          self._ad_dFdm_adj2_forms,
+                                          self._ad_d2Fdm2_adj_forms,
+                                          self._ad_d2Fdudm_forms,
+
+                                          ad_block_tag=self.ad_block_tag)
+
+                for dep in self._ad_dependencies_to_add:
+                    block.add_dependency(dep, no_duplicates=True)
+                # mesh = self._ad_problem.u.function_space().mesh()
+                # block.add_dependency(mesh, no_duplicates=True)
+
+                get_working_tape().add_block(block)
+
+            with stop_annotating():
+                out = solve(self, **kwargs)
+
+            if annotate:
+                block.add_output(self._ad_problem._ad_u.create_block_variable())
+
+            return out
+
+        return wrapper
+
+    @staticmethod
+    def _ad_annotate_solve_old(solve):
         @wraps(solve)
         def wrapper(self, **kwargs):
             """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
diff --git a/tests/firedrake/adjoint/test_assemble.py b/tests/firedrake/adjoint/test_assemble.py
index 22bd7540a1..5e891619a6 100644
--- a/tests/firedrake/adjoint/test_assemble.py
+++ b/tests/firedrake/adjoint/test_assemble.py
@@ -89,9 +89,11 @@ def test_assemble_1_forms_tlm(rg):
 
     h = rg.uniform(V)
     g = f.copy(deepcopy=True)
-    f.block_variable.tlm_value = h
-    tape.evaluate_tlm()
-    assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9)
+    Jhat(g)
+    assert (taylor_test(Jhat, g, h, dJdm=Jhat.tlm(h)) > 1.9)
+    # f.block_variable.tlm_value = h
+    # tape.evaluate_tlm()
+    # assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9)
 
 
 @pytest.mark.skipcomplex
diff --git a/tests/firedrake/adjoint/test_hessian.py b/tests/firedrake/adjoint/test_hessian.py
index e3ec975c97..f6942d9cbe 100644
--- a/tests/firedrake/adjoint/test_hessian.py
+++ b/tests/firedrake/adjoint/test_hessian.py
@@ -124,7 +124,6 @@ def test_function(rg):
     J = assemble(c ** 2 * u ** 2 * dx)
 
     Jhat = ReducedFunctional(J, [control_c, control_f])
-    dJdc, dJdf = compute_derivative(J, [control_c, control_f], apply_riesz=True)
 
     # Step direction for derivatives and convergence test
     h_c = Function(R, val=1.0)
@@ -161,6 +160,7 @@ def test_nonlinear(rg):
     Jhat = ReducedFunctional(J, Control(f))
 
     h = rg.uniform(V, 0, 10)
+    g = f.copy(deepcopy=True)
 
     J.block_variable.adj_value = 1.0
     f.block_variable.tlm_value = h
@@ -171,8 +171,6 @@ def test_nonlinear(rg):
     J.block_variable.hessian_value = 0
     tape.evaluate_hessian()
 
-    g = f.copy(deepcopy=True)
-
     dJdm = J.block_variable.tlm_value
     Hm = f.block_variable.hessian_value.dat.inner(h.dat)
     assert taylor_test(Jhat, g, h, dJdm=dJdm, Hm=Hm) > 2.8
@@ -223,25 +221,24 @@ def test_dirichlet(rg):
 def test_burgers(solve_type, rg):
     tape = Tape()
     set_working_tape(tape)
-    n = 100
-    mesh = UnitIntervalMesh(n)
-    V = FunctionSpace(mesh, "CG", 2)
+    nx = 50
+    nt = 5
+    mesh = UnitIntervalMesh(nx)
+    V = FunctionSpace(mesh, "CG", 1)
 
-    def Dt(u, u_, timestep):
-        return (u - u_)/timestep
+    def Dt(u, u_, dt):
+        return (u - u_)/dt
 
     x, = SpatialCoordinate(mesh)
-    pr = project(sin(2*pi*x), V, annotate=False)
-    ic = Function(V).assign(pr)
+    ic = Function(V).project(sin(2*pi*x))
 
     u_ = Function(V).assign(ic)
     u = Function(V).assign(ic)
     v = TestFunction(V)
 
-    nu = Constant(0.0001)
+    nu = Constant(1/100)
 
-    dt = 0.01
-    nt = 20
+    dt = Constant(1/nx)
 
     params = {
         'snes_rtol': 1e-10,
@@ -266,12 +263,6 @@ def Dt(u, u_, timestep):
             NonlinearVariationalProblem(F, u),
             solver_parameters=params)
 
-    if use_nlvs:
-        solver.solve()
-    else:
-        solve(F == 0, u, bc, solver_parameters=params)
-    u_.assign(u)
-
     for _ in range(nt):
         if use_nlvs:
             solver.solve()
@@ -282,16 +273,13 @@ def Dt(u, u_, timestep):
     J = assemble(u_*u_*dx + ic*ic*dx)
 
     Jhat = ReducedFunctional(J, Control(ic))
+
     h = rg.uniform(V)
     g = ic.copy(deepcopy=True)
-    J.block_variable.adj_value = 1.0
-    ic.block_variable.tlm_value = h
-    tape.evaluate_adj()
-    tape.evaluate_tlm()
-
-    J.block_variable.hessian_value = 0
-    tape.evaluate_hessian()
 
-    dJdm = J.block_variable.tlm_value
-    Hm = ic.block_variable.hessian_value.dat.inner(h.dat)
-    assert taylor_test(Jhat, g, h, dJdm=dJdm, Hm=Hm) > 2.9
+    taylor = taylor_to_dict(Jhat, g, h)
+    from pprint import pprint
+    pprint(taylor)
+    assert min(taylor['R0']['Rate']) > 0.95, taylor['R0']
+    assert min(taylor['R1']['Rate']) > 1.95, taylor['R1']
+    assert min(taylor['R2']['Rate']) > 2.95, taylor['R2']
diff --git a/tests/firedrake/adjoint/test_nlvs.py b/tests/firedrake/adjoint/test_nlvs.py
new file mode 100644
index 0000000000..c4f3f93b41
--- /dev/null
+++ b/tests/firedrake/adjoint/test_nlvs.py
@@ -0,0 +1,191 @@
+import pytest
+
+from firedrake import *
+from firedrake.adjoint import *
+
+
+@pytest.fixture(autouse=True)
+def handle_taping():
+    yield
+    tape = get_working_tape()
+    tape.clear_tape()
+
+
+@pytest.fixture(autouse=True, scope="module")
+def handle_annotation():
+    if not annotate_tape():
+        continue_annotation()
+    yield
+    # Ensure annotation is paused when we finish.
+    if annotate_tape():
+        pause_annotation()
+
+
+def forward(ic, dt, nt, bc_arg=None):
+    """Burgers equation solver."""
+    V = ic.function_space()
+
+    if bc_arg:
+        bc_val = bc_arg.copy(deepcopy=True)
+        bcs = [DirichletBC(V, bc_val, 1),
+               DirichletBC(V, 0, 2)]
+    else:
+        bcs = None
+
+    nu = Function(dt.function_space()).assign(0.1)
+
+    u0 = Function(V)
+    u1 = Function(V)
+    v = TestFunction(V)
+
+    F = ((u1 - u0)*v
+         + dt*u1*u1.dx(0)*v
+         + dt*nu*u1.dx(0)*v.dx(0))*dx
+
+    problem = NonlinearVariationalProblem(F, u1, bcs=bcs)
+    solver = NonlinearVariationalSolver(problem)
+
+    u1.assign(ic)
+
+    for i in range(nt):
+        u0.assign(u1)
+        solver.solve()
+        nu += dt
+        # if bc_arg:
+        #     bc_val.assign(bc_val + dt/nt)
+
+    J = assemble(u1*u1*dx)
+    return J
+
+
+@pytest.mark.skipcomplex
+@pytest.mark.parametrize("control_type", ["ic_control",
+                                          "dt_control",
+                                          "bc_control"])
+@pytest.mark.parametrize("bc_type", ["neumann_bc",
+                                     "dirichlet_bc"])
+def test_nlvs_adjoint(control_type, bc_type):
+    if control_type == 'bc_control' and bc_type == 'neumann_bc':
+        pytest.skip("Cannot use Neumann BCs as control")
+
+    nx = 100000
+    nt = 50
+
+    mesh = UnitIntervalMesh(nx)
+    x, = SpatialCoordinate(mesh)
+
+    V = FunctionSpace(mesh, "CG", 1)
+    R = FunctionSpace(mesh, "R", 0)
+
+    dt = Function(R).assign(1/nx)
+    ic = Function(V).interpolate(cos(2*pi*x))
+
+    dt0 = dt.copy(deepcopy=True)
+    ic0 = ic.copy(deepcopy=True)
+
+    if bc_type == 'neumann_bc':
+        bc_arg = None
+        bc_arg0 = None
+    elif bc_type == 'dirichlet_bc':
+        bc_arg = Function(R).assign(1.)
+        bc_arg0 = bc_arg.copy(deepcopy=True)
+    else:
+        raise ValueError(f"Unrecognised {bc_type=}")
+
+    if control_type == 'ic_control':
+        control = ic0
+    elif control_type == 'dt_control':
+        control = dt0
+    elif control_type == 'bc_control':
+        control = bc_arg0
+    else:
+        raise ValueError(f"Unrecognised {control_type=}")
+
+    PETSc.Sys.Print("record tape")
+    continue_annotation()
+    with set_working_tape() as tape:
+        J = forward(ic0, dt0, nt, bc_arg=bc_arg0)
+        Jhat = ReducedFunctional(J, Control(control), tape=tape)
+    pause_annotation()
+
+    if control_type == 'ic_control':
+        m = Function(V).assign(0.5*ic)
+        h = Function(V).interpolate(-0.5*cos(4*pi*x))
+
+        ic2 = m.copy(deepcopy=True)
+        dt2 = dt
+        bc_arg2 = bc_arg
+
+    elif control_type == 'dt_control':
+        m = Function(R).assign(0.05)
+        h = Function(R).assign(0.01)
+
+        ic2 = ic
+        dt2 = m.copy(deepcopy=True)
+        bc_arg2 = bc_arg
+
+    elif control_type == 'bc_control':
+        m = Function(R).assign(0.5)
+        h = Function(R).assign(-0.1)
+
+        ic2 = ic
+        dt2 = dt
+        bc_arg2 = m.copy(deepcopy=True)
+
+    from mpi4py import MPI
+
+    # # recompute component
+    # PETSc.Sys.Print("recompute test")
+    # assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2)) < 1e-14
+
+    # # tlm
+    # PETSc.Sys.Print("tlm test")
+    # Jhat(m)
+    # assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95
+
+    # # adjoint
+    # PETSc.Sys.Print("adjoint test")
+    # Jhat(m)
+    # assert taylor_test(Jhat, m, h) > 1.95
+
+    # # hessian
+    # PETSc.Sys.Print("hessian test")
+    # Jhat(m)
+    # taylor = taylor_to_dict(Jhat, m, h)
+    # from pprint import pprint
+    # pprint(taylor)
+
+    # assert min(taylor['R0']['Rate']) > 0.95
+    # assert min(taylor['R1']['Rate']) > 1.95
+    # assert min(taylor['R2']['Rate']) > 2.95
+
+    for _ in range(3):
+        stime = MPI.Wtime()
+        Jhat(m)
+        etime = MPI.Wtime()
+        PETSc.Sys.Print(f"Recompute time: {etime - stime:.4f}")
+
+    for _ in range(3):
+        stime = MPI.Wtime()
+        Jhat.derivative()
+        etime = MPI.Wtime()
+        PETSc.Sys.Print(f"Derivative time: {etime - stime:.4f}")
+
+    for _ in range(3):
+        stime = MPI.Wtime()
+        Jhat.tlm(h)
+        etime = MPI.Wtime()
+        PETSc.Sys.Print(f"TLM time: {etime - stime:.4f}")
+
+    for _ in range(3):
+        stime = MPI.Wtime()
+        Jhat.hessian(h, evaluate_tlm=False)
+        etime = MPI.Wtime()
+        PETSc.Sys.Print(f"Hessian time: {etime - stime:.4f}")
+
+
+if __name__ == "__main__":
+    control_type = "ic_control"
+    bc_type = "neumann_bc"
+    PETSc.Sys.Print(f"{control_type=} | {bc_type=}")
+    test_nlvs_adjoint(control_type, bc_type)

From 1128698646822bed716977dbc039528a8aea3582 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Mon, 8 Jun 2026 17:26:17 +0100
Subject: [PATCH 2/8] Disable annotation when creating solver caches

---
 firedrake/adjoint_utils/variational_solver.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index b5afc05733..88de125258 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -368,10 +368,11 @@ def wrapper(self, **kwargs):
                         "MissingMathsError: we do not know how to differentiate through a variational inequality")
 
                 if len(self._ad_solver_cache) == 0:
-                    self._ad_cache_forward_solver()
-                    self._ad_cache_tlm_solver()
-                    self._ad_cache_adj_solver()
-                    self._ad_cache_hessian_solver()
+                    with stop_annotating():
+                        self._ad_cache_forward_solver()
+                        self._ad_cache_tlm_solver()
+                        self._ad_cache_adj_solver()
+                        self._ad_cache_hessian_solver()
 
                 block = CachedSolverBlock(self._ad_problem.u,
                                           self._ad_bcs,

From e32192a3723ada702ec0e731e98b23f82001d9f0 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Mon, 8 Jun 2026 22:36:42 +0100
Subject: [PATCH 3/8] Remove adj_sol_buf

A bit of a doubling up on the already-present adj_sol on
solver blocks.
---
 firedrake/adjoint_utils/blocks/solving.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 2c75fe9cb3..6f791b2192 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -66,7 +66,6 @@ def __init__(self, func, bcs, cached_solvers,
         self.adj_residual = adj_residual
 
         self.adj_sol = adj_sol
-        self.adj_sol_buf = adj_sol.copy(deepcopy=True)
         self.adj2_sol = adj2_sol
         self.tlm_output = tlm_output
         self.d2Fdu2_form = d2Fdu2_form
@@ -190,7 +189,6 @@ def solve_adj_equation(self, rhs, compute_boundary):
         solver.solve()
 
         self.adj_sol.assign(adj_sol)
-        self.adj_sol_buf.assign(adj_sol)
 
         if compute_boundary:
             adj_sol_bc = firedrake.assemble(self.adj_residual)
@@ -214,7 +212,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         # self.adj_sol is shared between all blocks that this NLVS
         # generates so we can't store it there. Instead store it
         # in self.adj_sol_buf which is owned by this block only.
-        self.adj_sol_buf.assign(adj_sol)
+        self.adj_sol.assign(adj_sol)
 
         prepared = {
             "adj_sol": adj_sol.copy(deepcopy=True),
@@ -239,7 +237,6 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
         return dFdm
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
-        self.adj_sol.assign(self.adj_sol_buf)
         self.update_dependencies(use_output=True)
         self.update_hessian_dependencies()
 

From db9f4228e28c0651d5580a9bc1f3d7565a65aae5 Mon Sep 17 00:00:00 2001
From: Josh Hope-Collins <joshua.hope-collins13@imperial.ac.uk>
Date: Thu, 11 Jun 2026 13:09:18 +0100
Subject: [PATCH 4/8] Gather adjoint solver cached data together (#5161)

* reduce nlvs adjoint test size

* nlvs adjoint caching - bundle up recompute objects
---
 firedrake/adjoint_utils/blocks/solving.py     | 128 +++++------
 firedrake/adjoint_utils/variational_solver.py | 214 +++++++++++-------
 tests/firedrake/adjoint/test_nlvs.py          |  56 ++---
 3 files changed, 212 insertions(+), 186 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 6f791b2192..f3a84f58a1 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -40,48 +40,29 @@ class SolverType(Enum):
 
 
 class CachedSolverBlock(Block):
-    def __init__(self, func, bcs, cached_solvers,
-                 is_linear, replaced_dependencies,
-                 tlm_rhs, replaced_tlms, tlm_dFdm_forms,
-                 adj_rhs, adj_dFdm_forms, adj_residual,
-                 adj_sol, adj2_sol, tlm_output,
-                 d2Fdu2_form, d2Fdmdu_forms,
-                 dFdm_adj2_forms, d2Fdm2_adj_forms,
-                 d2Fdudm_forms,
+    def __init__(self, forward_cache, tangent_cache, adjoint_cache, hessian_cache,
                  ad_block_tag=None):
         super().__init__(ad_block_tag=ad_block_tag)
 
-        self.func = func
-        self.bcs = bcs
-        self.cached_solvers = cached_solvers
-        self.replaced_dependencies = replaced_dependencies
-        self.is_linear = is_linear
-
-        self.tlm_rhs = tlm_rhs
-        self.replaced_tlms = replaced_tlms
-        self.tlm_dFdm_forms = tlm_dFdm_forms
+        self.forward_cache = forward_cache
+        self.tangent_cache = tangent_cache
+        self.adjoint_cache = adjoint_cache
+        self.hessian_cache = hessian_cache
+        self.is_linear = forward_cache.is_linear
 
-        self.adj_rhs = adj_rhs
-        self.adj_dFdm_forms = adj_dFdm_forms
-        self.adj_residual = adj_residual
-
-        self.adj_sol = adj_sol
-        self.adj2_sol = adj2_sol
-        self.tlm_output = tlm_output
-        self.d2Fdu2_form = d2Fdu2_form
-        self.d2Fdmdu_forms = d2Fdmdu_forms
-        self.dFdm_adj2_forms = dFdm_adj2_forms
-        self.d2Fdm2_adj_forms = d2Fdm2_adj_forms
-        self.d2Fdudm_forms = d2Fdudm_forms
+        # The adj_sol in the cached forms is shared by all blocks.
+        # This adj_sol belongs to this block specifically so we can
+        # stash the adjoint solution for the hessian calculation.
+        self.adj_sol = adjoint_cache.adj_sol.copy(deepcopy=True, annotate=False)
 
     def _coefficient_dependencies(self, dependencies=None):
         dependencies = dependencies or self.get_dependencies()
-        return dependencies[:len(self.replaced_dependencies)]
+        return dependencies[:len(self.forward_cache.replaced_deps)]
 
     def _bc_dependencies(self, dependencies=None):
         dependencies = dependencies or self.get_dependencies()
-        if len(self.bcs) > 0:
-            return dependencies[-len(self.bcs):]
+        if len(self.forward_cache.bcs) > 0:
+            return dependencies[-len(self.forward_cache.bcs):]
         else:
             return []
 
@@ -90,7 +71,7 @@ def update_dependencies(self, use_output=False):
         """
         # Update the coefficients in the form.
         # Use the fact that zip will use the shorter length.
-        for replaced_dep, dep in zip(self.replaced_dependencies,
+        for replaced_dep, dep in zip(self.forward_cache.replaced_deps,
                                      self._coefficient_dependencies()):
             replaced_dep.assign(dep.saved_output)
 
@@ -101,24 +82,24 @@ def update_dependencies(self, use_output=False):
         # Jacobian is correct.
         if use_output:
             output = self.get_outputs()[0].saved_output
-            self.cached_solvers[FORWARD]._problem.u.assign(output)
+            self.forward_cache.solver._problem.u.assign(output)
 
         # Update the boundary conditions
-        for replaced_dep, dep in zip(self.bcs, self._bc_dependencies()):
+        for replaced_dep, dep in zip(self.forward_cache.bcs, self._bc_dependencies()):
             replaced_dep.set_value(dep.saved_output.function_arg)
 
     def update_tlm_dependencies(self):
         """Update all dependencies of the tlm solve.
         """
-        for replaced_dep, dep in zip(self.replaced_tlms,
+        for replaced_dep, dep in zip(self.tangent_cache.replaced_tlms,
                                      self._coefficient_dependencies()):
-            if dep.output == self.func and not self.is_linear:
+            if dep.output == self.forward_cache.func and not self.is_linear:
                 continue
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
             replaced_dep.assign(dep.tlm_value)
 
-        for replaced_dep, dep in zip(self.bcs, self._bc_dependencies()):
+        for replaced_dep, dep in zip(self.forward_cache.bcs, self._bc_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 bc_val = 0
             else:
@@ -132,6 +113,9 @@ def update_adj_dependencies(self):
     def update_hessian_dependencies(self):
         # TODO: Anything else to do here?
         self.update_tlm_dependencies()
+        # update the adj_sol in the cached forms with
+        # the adj_sol value owned by this block.
+        self.hessian_cache.adj_sol.assign(self.adj_sol)
 
     def _compute_boundary(self, relevant_dependencies):
         return any(isinstance(dep.output, firedrake.DirichletBC)
@@ -143,7 +127,7 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
     def recompute_component(self, inputs, block_variable, idx, prepared):
         self.update_dependencies(use_output=False)
 
-        solver = self.cached_solvers[FORWARD]
+        solver = self.forward_cache.solver
         solver.solve()
         result = solver._problem.u.copy(deepcopy=True)
 
@@ -161,42 +145,40 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         self.update_tlm_dependencies()
 
         # Assemble the rhs of (dF/du)(du/dm) = -dF/dm
-        self.tlm_rhs.zero()
-        for dFdm, dep in zip(self.tlm_dFdm_forms, self.get_dependencies()):
+        tlm_rhs = self.tangent_cache.rhs
+        tlm_rhs.zero()
+        for dFdm, dep in zip(self.tangent_cache.dFdm_forms, self.get_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
-            if dep.output is self.func and not self.is_linear:  # Can't compute dependence on initial guess
+            if dep.output is self.forward_cache.func and not self.is_linear:  # Can't compute dependence on initial guess
                 continue
-            self.tlm_rhs += firedrake.assemble(dFdm)
+            tlm_rhs += firedrake.assemble(dFdm)
 
         # Solve for dudm
-        solver = self.cached_solvers[TLM]
+        solver = self.tangent_cache.solver
         solver._problem.u.zero()
         solver.solve()
         result = solver._problem.u.copy(deepcopy=True)
         return result
 
     def solve_adj_equation(self, rhs, compute_boundary):
-        for bc in self.bcs:
+        for bc in self.forward_cache.bcs:
             bc.homogenize()
 
-        solver = self.cached_solvers[ADJOINT]
-        adj_sol = solver._problem.u
+        adj_rhs = self.adjoint_cache.rhs
+        adj_sol = self.adjoint_cache.adj_sol
 
-        self.adj_rhs.assign(rhs)
+        adj_rhs.assign(rhs)
         adj_sol.zero()
-
-        solver.solve()
-
-        self.adj_sol.assign(adj_sol)
+        self.adjoint_cache.solver.solve()
 
         if compute_boundary:
-            adj_sol_bc = firedrake.assemble(self.adj_residual)
+            adj_sol_bc = firedrake.assemble(self.adjoint_cache.residual)
             adj_sol_bc = adj_sol_bc.riesz_representation("l2")
         else:
             adj_sol_bc = None
 
-        return adj_sol, adj_sol_bc
+        return adj_sol.copy(deepcopy=True), adj_sol_bc
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         self.update_dependencies(use_output=True)
@@ -208,20 +190,22 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
 
         adj_sol, adj_sol_bc = self.solve_adj_equation(dJdu, compute_boundary)
 
-        # store adj_sol for Hessian computation later.
-        # self.adj_sol is shared between all blocks that this NLVS
-        # generates so we can't store it there. Instead store it
-        # in self.adj_sol_buf which is owned by this block only.
+        # store adj_sol for Hessian computation later, or for inspecting
+        # adjoint sensitivities etc.
+        # self.adj_sol is owned by this block, whereas self.hessian_cache.adj_sol
+        # is shared between all blocks that the NLVS generates because it is the
+        # one in the cached forms, so we will update it as necessary if/when each
+        # block calculates the Hessian action.
         self.adj_sol.assign(adj_sol)
 
         prepared = {
-            "adj_sol": adj_sol.copy(deepcopy=True),
+            "adj_sol": adj_sol,
             "adj_sol_bc": adj_sol_bc
         }
         return prepared
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
-        if block_variable.output == self.func and not self.is_linear:
+        if block_variable.output == self.forward_cache.func and not self.is_linear:
             return None
 
         if isinstance(block_variable.output, firedrake.DirichletBC):
@@ -232,7 +216,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
             )
 
         # assemble sensititivy comment
-        dFdm = firedrake.assemble(self.adj_dFdm_forms[idx])
+        dFdm = firedrake.assemble(self.adjoint_cache.dFdm_forms[idx])
 
         return dFdm
 
@@ -254,15 +238,15 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
         hessian_rhs = hessian_input.copy(deepcopy=True)
 
         # tlm_output contribution
-        self.tlm_output.assign(tlm_output)
-        hessian_rhs -= firedrake.assemble(self.d2Fdu2_form)
+        self.hessian_cache.tlm_output.assign(tlm_output)
+        hessian_rhs -= firedrake.assemble(self.hessian_cache.d2Fdu2_form)
 
         # tlm_input contribution
-        for d2Fdmdu, dep in zip(self.d2Fdmdu_forms,
+        for d2Fdmdu, dep in zip(self.hessian_cache.d2Fdmdu_forms,
                                 self._coefficient_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
-            if dep.output is self.func and not self.is_linear:  # Can't compute dependence on initial guess
+            if dep.output is self.forward_cache.func and not self.is_linear:  # Can't compute dependence on initial guess
                 continue
             if len(d2Fdmdu.integrals()) > 0:
                 hessian_rhs -= firedrake.assemble(d2Fdmdu)
@@ -271,10 +255,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
         compute_boundary = self._compute_boundary(relevant_dependencies)
         adj2_sol, adj2_sol_bc = self.solve_adj_equation(hessian_rhs, compute_boundary)
 
-        self.adj2_sol.assign(adj2_sol)
+        self.hessian_cache.adj2_sol.assign(adj2_sol)
 
         prepared = {
-            "adj2_sol": adj2_sol.copy(deepcopy=True),
+            "adj2_sol": adj2_sol,
             "adj2_sol_bc": adj2_sol_bc,
         }
 
@@ -283,7 +267,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None):
         m = block_variable.output
 
-        if m is self.func and not self.is_linear:
+        if m is self.forward_cache.func and not self.is_linear:
             return None
 
         if isinstance(m, firedrake.DirichletBC):
@@ -299,14 +283,14 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
                 continue
             if dep.tlm_value is None:
                 continue
-            if dep.output is self.func and not self.is_linear:
+            if dep.output is self.forward_cache.func and not self.is_linear:
                 continue
-            relevant_d2Fdm2_forms.append(self.d2Fdm2_adj_forms[idx][i])
+            relevant_d2Fdm2_forms.append(self.hessian_cache.d2Fdm2_adj_forms[idx][i])
 
         hessian_output = 0
 
-        for form in (self.d2Fdudm_forms[idx],
-                     self.dFdm_adj2_forms[idx],
+        for form in (self.hessian_cache.d2Fdudm_forms[idx],
+                     self.hessian_cache.dFdm_adj2_forms[idx],
                      *relevant_d2Fdm2_forms):
             if not form.empty():
                 hessian_output += firedrake.assemble(-form)
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 88de125258..814748fa3a 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -1,5 +1,5 @@
 import copy
-from functools import wraps
+from functools import wraps, cached_property
 from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations
 from firedrake.adjoint_utils.blocks import (
     NonlinearVariationalSolveBlock, CachedSolverBlock)
@@ -8,6 +8,57 @@
 from ufl import replace, Action
 from ufl.algorithms import expand_derivatives
 from types import SimpleNamespace
+from collections import namedtuple
+
+
+ForwardSolveRecomputeCache = namedtuple(
+    'ForwardSolveRecomputeCache',
+    field_names=[
+        "func",
+        "bcs",
+        "solver",
+        "is_linear",
+        "replaced_deps",
+    ]
+)
+
+AdjointSolveRecomputeCache = namedtuple(
+    'AdjointSolveRecomputeCache',
+    field_names=[
+        "adj_sol",
+        "rhs",
+        "dFdm_forms",
+        "residual",
+        "solver",
+        "dFdu_adj",
+    ]
+)
+
+TangentSolveRecomputeCache = namedtuple(
+    "TangentSolveRecomputeCache",
+    field_names=[
+        "tlm_val",
+        "rhs",
+        "dFdm_forms",
+        "solver",
+        "replaced_tlms",
+        "dFdu",
+    ]
+)
+
+HessianSolveRecomputeCache = namedtuple(
+    "HessianSolveRecomputeCache",
+    field_names=[
+        "adj_sol",
+        "adj2_sol",
+        "tlm_output",
+        "d2Fdu2_form",
+        "d2Fdmdu_forms",
+        "dFdm_adj2_forms",
+        "d2Fdm2_adj_forms",
+        "d2Fdudm_forms",
+    ]
+)
 
 
 class NonlinearVariationalProblemMixin:
@@ -73,12 +124,13 @@ def wrapper(self, problem, *args, **kwargs):
 
         return wrapper
 
-    def _ad_cache_forward_solver(self):
+    @cached_property
+    @no_annotations
+    def _ad_forward_cache(self):
         from firedrake import (
             DirichletBC,
             NonlinearVariationalProblem,
             NonlinearVariationalSolver)
-        from firedrake.adjoint_utils.blocks.solving import FORWARD
 
         problem = self._ad_problem
 
@@ -125,20 +177,26 @@ def _ad_cache_forward_solver(self):
         # The block need handles to the newly created
         # objects to update their values when recomputing.
         self._ad_dependencies_to_add = (*replace_map.keys(), *bcs)
-        self._ad_replaced_dependencies = tuple(replace_map.values())
-        self._ad_bcs = bcs_new
-        self._ad_solver_cache[FORWARD] = nlvs
 
-    def _ad_cache_tlm_solver(self):
+        return ForwardSolveRecomputeCache(
+            func=self._ad_problem.u,
+            bcs=bcs_new,
+            solver=nlvs,
+            is_linear=self._ad_problem.is_linear,
+            replaced_deps=tuple(replace_map.values()),
+        )
+
+    @cached_property
+    @no_annotations
+    def _ad_tangent_cache(self):
         from firedrake import (
             Function, Cofunction, derivative, TrialFunction,
             LinearVariationalProblem, LinearVariationalSolver)
-        from firedrake.adjoint_utils.blocks.solving import FORWARD, TLM
 
         # If we build the TLM form from the cached
         # forward solve form then we can update exactly
         # the same coefficients/boundary conditions.
-        nlvp = self._ad_solver_cache[FORWARD]._problem
+        nlvp = self._ad_forward_cache.solver._problem
 
         F = nlvp.F
         u = nlvp.u
@@ -153,25 +211,22 @@ def _ad_cache_tlm_solver(self):
         dFdm = Cofunction(V.dual())
         dudm = Function(V)
 
-        self._ad_dFdu = dFdu
-
         # Reuse the same bcs as the forward problem.
         # TODO: Think about if we should use new bcs.
         # TODO: solver_parameters
-        lvp = LinearVariationalProblem(dFdu, dFdm, dudm, bcs=self._ad_bcs)
+        lvp = LinearVariationalProblem(dFdu, dFdm, dudm, bcs=self._ad_forward_cache.bcs)
         lvs = LinearVariationalSolver(
             lvp,
             *self._ad_args_kwargs.tlm_args,
             **self._ad_args_kwargs.tlm_kwargs)
 
-        self._ad_solver_cache[TLM] = lvs
-        self._ad_tlm_rhs = dFdm
+        tlm_rhs = dFdm
 
         # Do all the symbolic work for calculating dF/dm up front
         # so we only pay for the numeric calculations at run time.
         replaced_tlms = []
         dFdm_tlm_forms = []
-        for m in self._ad_replaced_dependencies:
+        for m in self._ad_forward_cache.replaced_deps:
             mtlm = m.copy(deepcopy=True)
             replaced_tlms.append(mtlm)
 
@@ -180,21 +235,28 @@ def _ad_cache_tlm_solver(self):
             dFdm = expand_derivatives(dFdm)
             dFdm_tlm_forms.append(dFdm)
 
-        # We'll need to update the replaced_tlm
-        # values and assemble the dFdm forms
-        self._ad_tlm_dFdm_forms = dFdm_tlm_forms
-        self._ad_replaced_tlms = replaced_tlms
+        tlm_val = Function(V)
+
+        return TangentSolveRecomputeCache(
+            tlm_val=tlm_val,
+            rhs=tlm_rhs,
+            dFdm_forms=dFdm_tlm_forms,
+            solver=lvs,
+            replaced_tlms=replaced_tlms,
+            dFdu=dFdu,
+        )
 
-    def _ad_cache_adj_solver(self):
+    @cached_property
+    @no_annotations
+    def _ad_adjoint_cache(self):
         from firedrake import (
             Function, Cofunction, TrialFunction, Argument,
             LinearVariationalProblem, LinearVariationalSolver)
-        from firedrake.adjoint_utils.blocks.solving import FORWARD, ADJOINT
 
         # If we build the adjoint form from the cached
         # forward solve form then we can update exactly
         # the same coefficients/boundary conditions.
-        nlvp = self._ad_solver_cache[FORWARD]._problem
+        nlvp = self._ad_forward_cache.solver._problem
 
         F = nlvp.F
         u = nlvp.u
@@ -206,7 +268,7 @@ def _ad_cache_adj_solver(self):
         # Then for the _partial_ derivatives:
         # (dF/du)*(du/dm) + dF/dm = 0 so we calculate:
         # (dF/du)*(du/dm) = -dF/dm
-        dFdu = self._ad_dFdu
+        dFdu = self._ad_tangent_cache.dFdu
         try:
             dFdu_adj = adjoint(dFdu)
         except ValueError:
@@ -214,8 +276,6 @@ def _ad_cache_adj_solver(self):
             # as dFdu might have been simplied to an empty Form
             dFdu_adj = adjoint(dFdu, derivatives_expanded=True)
 
-        self._ad_dFdu_adj = dFdu_adj
-
         # This will be the rhs of the adjoint problem
         dJdu = Cofunction(V.dual())
         adj_sol = Function(V)
@@ -223,19 +283,18 @@ def _ad_cache_adj_solver(self):
         # Reuse the same bcs as the forward problem.
         # TODO: Think about if we should use new bcs.
         # TODO: solver_parameters
-        lvp = LinearVariationalProblem(dFdu_adj, dJdu, adj_sol, bcs=self._ad_bcs)
+        lvp = LinearVariationalProblem(dFdu_adj, dJdu, adj_sol, bcs=self._ad_forward_cache.bcs)
         lvs = LinearVariationalSolver(
             lvp,
             *self._ad_args_kwargs.adj_args,
             **self._ad_args_kwargs.adj_kwargs)
 
-        self._ad_solver_cache[ADJOINT] = lvs
-        self._ad_adj_rhs = dJdu
-
+        # We'll need to assemble these forms to calculate
+        # the adj_component for each dependency.
         # Do all the symbolic work for calculating dJ/du up front
         # so we only pay for the numeric calculations at run time.
         dFdm_adj_forms = []
-        for m in self._ad_replaced_dependencies:
+        for m in self._ad_forward_cache.replaced_deps:
             # Action of adjoint solution on dFdm
             # TODO: Which of the two implementations should we use?
             dFdm = derivative(-F, m, TrialFunction(m.function_space()))
@@ -262,18 +321,24 @@ def _ad_cache_adj_solver(self):
         # we'll need the residual of the adjoint equation without
         # any DirichletBC using the solution calculated with
         # homogeneous DirichletBCs.
-        self._ad_adj_residual = dJdu - action(dFdu_adj, adj_sol)
-
-        # We'll need to assemble these forms to calculate
-        # the adj_component for each dependency.
-        self._ad_adj_dFdm_forms = dFdm_adj_forms
-
-    def _ad_cache_hessian_solver(self):
+        adj_residual = dJdu - action(dFdu_adj, adj_sol)
+
+        return AdjointSolveRecomputeCache(
+            adj_sol=adj_sol,
+            rhs=dJdu,
+            dFdm_forms=dFdm_adj_forms,
+            residual=adj_residual,
+            solver=lvs,
+            dFdu_adj=dFdu_adj,
+        )
+
+    @cached_property
+    @no_annotations
+    def _ad_hessian_cache(self):
         from firedrake import (
             Function, TestFunction)
-        from firedrake.adjoint_utils.blocks.solving import FORWARD
 
-        nlvp = self._ad_solver_cache[FORWARD]._problem
+        nlvp = self._ad_forward_cache.solver._problem
         F = nlvp.F
         u = nlvp.u
         V = u.function_space()
@@ -282,7 +347,7 @@ def _ad_cache_hessian_solver(self):
 
         # Calculate d^2F/du^2 * du/dm * dm
         # where dm is direction for tlm action so du/dm * dm is tlm output
-        dFdu = self._ad_dFdu
+        dFdu = self._ad_tangent_cache.dFdu
         tlm_output = Function(V)
         d2Fdu2 = derivative(dFdu, u, tlm_output)
         # print()
@@ -292,33 +357,26 @@ def _ad_cache_hessian_solver(self):
         # print()
         d2Fdu2 = expand_derivatives(d2Fdu2)
 
-        self._ad_tlm_output = tlm_output
-
         adj_sol = Function(V)
-        self._ad_adj_sol = adj_sol
 
         # Contribution from tlm_output
         if len(d2Fdu2.integrals()) > 0:
             d2Fdu2_form = action(adjoint(d2Fdu2), adj_sol)
         else:
             d2Fdu2_form = d2Fdu2
-        self._ad_d2Fdu2_form = d2Fdu2_form
 
         # Contributions from each tlm_input
-        dFdu_adj = action(self._ad_dFdu_adj, adj_sol)
+        dFdu_adj = action(self._ad_adjoint_cache.dFdu_adj, adj_sol)
         d2Fdmdu_forms = []
-        for m, dm in zip(self._ad_replaced_dependencies,
-                         self._ad_replaced_tlms):
+        for m, dm in zip(self._ad_forward_cache.replaced_deps,
+                         self._ad_tangent_cache.replaced_tlms):
             d2Fdmdu = expand_derivatives(
                 derivative(dFdu_adj, m, dm))
 
             d2Fdmdu_forms.append(d2Fdmdu)
 
-        self._ad_d2Fdmdu_forms = d2Fdmdu_forms
-
         # 2. Forms to calculate contribution from each control
         adj2_sol = Function(V)
-        self._ad_adj2_sol = adj2_sol
 
         Fadj = action(F, adj_sol)
         Fadj2 = action(F, adj2_sol)
@@ -326,7 +384,7 @@ def _ad_cache_hessian_solver(self):
         dFdm_adj2_forms = []
         d2Fdm2_adj_forms = []
         d2Fdudm_forms = []
-        for m in self._ad_replaced_dependencies:
+        for m in self._ad_forward_cache.replaced_deps:
             dm = TestFunction(m.function_space())
             dFdm_adj2 = expand_derivatives(
                 derivative(Fadj2, m, dm))
@@ -341,17 +399,24 @@ def _ad_cache_hessian_solver(self):
             d2Fdudm_forms.append(d2Fdudm)
 
             d2Fdm2_adj_forms_k = []
-            for m2, dm2 in zip(self._ad_replaced_dependencies,
-                               self._ad_replaced_tlms):
+            for m2, dm2 in zip(self._ad_forward_cache.replaced_deps,
+                               self._ad_tangent_cache.replaced_tlms):
                 d2Fdm2_adj = expand_derivatives(
                     derivative(dFdm_adj, m2, dm2))
                 d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
 
             d2Fdm2_adj_forms.append(d2Fdm2_adj_forms_k)
 
-        self._ad_dFdm_adj2_forms = dFdm_adj2_forms
-        self._ad_d2Fdm2_adj_forms = d2Fdm2_adj_forms
-        self._ad_d2Fdudm_forms = d2Fdudm_forms
+        return HessianSolveRecomputeCache(
+            adj_sol=adj_sol,
+            adj2_sol=adj2_sol,
+            tlm_output=tlm_output,
+            d2Fdu2_form=d2Fdu2_form,
+            d2Fdmdu_forms=d2Fdmdu_forms,
+            dFdm_adj2_forms=dFdm_adj2_forms,
+            d2Fdm2_adj_forms=d2Fdm2_adj_forms,
+            d2Fdudm_forms=d2Fdudm_forms,
+        )
 
     @staticmethod
     def _ad_annotate_solve(solve):
@@ -367,37 +432,12 @@ def wrapper(self, **kwargs):
                     raise ValueError(
                         "MissingMathsError: we do not know how to differentiate through a variational inequality")
 
-                if len(self._ad_solver_cache) == 0:
-                    with stop_annotating():
-                        self._ad_cache_forward_solver()
-                        self._ad_cache_tlm_solver()
-                        self._ad_cache_adj_solver()
-                        self._ad_cache_hessian_solver()
-
-                block = CachedSolverBlock(self._ad_problem.u,
-                                          self._ad_bcs,
-                                          self._ad_solver_cache,
-                                          self._ad_problem.is_linear,
-                                          self._ad_replaced_dependencies,
-
-                                          self._ad_tlm_rhs,
-                                          self._ad_replaced_tlms,
-                                          self._ad_tlm_dFdm_forms,
-
-                                          self._ad_adj_rhs,
-                                          self._ad_adj_dFdm_forms,
-                                          self._ad_adj_residual,
-
-                                          self._ad_adj_sol,
-                                          self._ad_adj2_sol,
-                                          self._ad_tlm_output,
-                                          self._ad_d2Fdu2_form,
-                                          self._ad_d2Fdmdu_forms,
-                                          self._ad_dFdm_adj2_forms,
-                                          self._ad_d2Fdm2_adj_forms,
-                                          self._ad_d2Fdudm_forms,
-
-                                          ad_block_tag=self.ad_block_tag)
+                block = CachedSolverBlock(
+                    self._ad_forward_cache,
+                    self._ad_tangent_cache,
+                    self._ad_adjoint_cache,
+                    self._ad_hessian_cache,
+                    ad_block_tag=self.ad_block_tag)
 
                 for dep in self._ad_dependencies_to_add:
                     block.add_dependency(dep, no_duplicates=True)
diff --git a/tests/firedrake/adjoint/test_nlvs.py b/tests/firedrake/adjoint/test_nlvs.py
index c4f3f93b41..9eb7c7a005 100644
--- a/tests/firedrake/adjoint/test_nlvs.py
+++ b/tests/firedrake/adjoint/test_nlvs.py
@@ -68,8 +68,8 @@ def test_nlvs_adjoint(control_type, bc_type):
     if control_type == 'bc_control' and bc_type == 'neumann_bc':
         pytest.skip("Cannot use Neumann BCs as control")
 
-    nx = 100000
-    nt = 50
+    nx = 100
+    nt = 5
 
     mesh = UnitIntervalMesh(nx)
     x, = SpatialCoordinate(mesh)
@@ -77,7 +77,7 @@ def test_nlvs_adjoint(control_type, bc_type):
     V = FunctionSpace(mesh, "CG", 1)
     R = FunctionSpace(mesh, "R", 0)
 
-    dt = Function(R).assign(1/nx)
+    dt = Function(R).assign(2/nx)
     ic = Function(V).interpolate(cos(2*pi*x))
 
     dt0 = dt.copy(deepcopy=True)
@@ -134,30 +134,32 @@ def test_nlvs_adjoint(control_type, bc_type):
 
     from mpi4py import MPI
 
-    # # recompute component
-    # PETSc.Sys.Print("recompute test")
-    # assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2)) < 1e-14
-
-    # # tlm
-    # PETSc.Sys.Print("tlm test")
-    # Jhat(m)
-    # assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95
-
-    # # adjoint
-    # PETSc.Sys.Print("adjoint test")
-    # Jhat(m)
-    # assert taylor_test(Jhat, m, h) > 1.95
-
-    # # hessian
-    # PETSc.Sys.Print("hessian test")
-    # Jhat(m)
-    # taylor = taylor_to_dict(Jhat, m, h)
-    # from pprint import pprint
-    # pprint(taylor)
-
-    # assert min(taylor['R0']['Rate']) > 0.95
-    # assert min(taylor['R1']['Rate']) > 1.95
-    # assert min(taylor['R2']['Rate']) > 2.95
+    # recompute component
+    PETSc.Sys.Print("recompute test")
+    assert abs(Jhat(m) - forward(ic2, dt2, nt, bc_arg=bc_arg2)) < 1e-14
+
+    # tlm
+    PETSc.Sys.Print("tlm test")
+    Jhat(m)
+    assert taylor_test(Jhat, m, h, dJdm=Jhat.tlm(h)) > 1.95
+
+    # adjoint
+    PETSc.Sys.Print("adjoint test")
+    Jhat(m)
+    assert taylor_test(Jhat, m, h) > 1.95
+
+    # hessian
+    PETSc.Sys.Print("hessian test")
+    Jhat(m)
+    taylor = taylor_to_dict(Jhat, m, h)
+    from pprint import pprint
+    pprint(taylor)
+
+    assert min(taylor['R0']['Rate']) > 0.95
+    assert min(taylor['R1']['Rate']) > 1.95
+    assert min(taylor['R2']['Rate']) > 2.95
+
+    return
 
     for _ in range(3):
         stime = MPI.Wtime()

From a9adaa408ad196981ddc35fdd2bc082eea756422 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Wed, 17 Jun 2026 03:24:57 +1000
Subject: [PATCH 5/8] Try to handle solver caching for linear forms (-> #4638)
 (#5158)

* Remove linearity checks from adjoint evaluation

* Cache J and Jp from NLVP on forward solver

Co-authored-by: Josh Hope-Collins <jth113@ic.ac.uk>
---
 firedrake/adjoint_utils/blocks/solving.py     | 16 ++---
 firedrake/adjoint_utils/variational_solver.py | 72 +++++++++----------
 2 files changed, 41 insertions(+), 47 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index f3a84f58a1..5739c6041a 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -48,7 +48,6 @@ def __init__(self, forward_cache, tangent_cache, adjoint_cache, hessian_cache,
         self.tangent_cache = tangent_cache
         self.adjoint_cache = adjoint_cache
         self.hessian_cache = hessian_cache
-        self.is_linear = forward_cache.is_linear
 
         # The adj_sol in the cached forms is shared by all blocks.
         # This adj_sol belongs to this block specifically so we can
@@ -93,7 +92,7 @@ def update_tlm_dependencies(self):
         """
         for replaced_dep, dep in zip(self.tangent_cache.replaced_tlms,
                                      self._coefficient_dependencies()):
-            if dep.output == self.forward_cache.func and not self.is_linear:
+            if dep.output == self.forward_cache.func:
                 continue
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
@@ -150,7 +149,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         for dFdm, dep in zip(self.tangent_cache.dFdm_forms, self.get_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
-            if dep.output is self.forward_cache.func and not self.is_linear:  # Can't compute dependence on initial guess
+            if dep.output is self.forward_cache.func:  # Can't compute dependence on initial guess
                 continue
             tlm_rhs += firedrake.assemble(dFdm)
 
@@ -205,7 +204,7 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         return prepared
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
-        if block_variable.output == self.forward_cache.func and not self.is_linear:
+        if block_variable.output == self.forward_cache.func:
             return None
 
         if isinstance(block_variable.output, firedrake.DirichletBC):
@@ -239,14 +238,15 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
 
         # tlm_output contribution
         self.hessian_cache.tlm_output.assign(tlm_output)
-        hessian_rhs -= firedrake.assemble(self.hessian_cache.d2Fdu2_form)
+        if not self.hessian_cache.d2Fdu2_form.empty():
+            hessian_rhs -= firedrake.assemble(self.hessian_cache.d2Fdu2_form)
 
         # tlm_input contribution
         for d2Fdmdu, dep in zip(self.hessian_cache.d2Fdmdu_forms,
                                 self._coefficient_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
-            if dep.output is self.forward_cache.func and not self.is_linear:  # Can't compute dependence on initial guess
+            if dep.output is self.forward_cache.func:  # Can't compute dependence on initial guess
                 continue
             if len(d2Fdmdu.integrals()) > 0:
                 hessian_rhs -= firedrake.assemble(d2Fdmdu)
@@ -267,7 +267,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
     def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None):
         m = block_variable.output
 
-        if m is self.forward_cache.func and not self.is_linear:
+        if m is self.forward_cache.func:
             return None
 
         if isinstance(m, firedrake.DirichletBC):
@@ -283,7 +283,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
                 continue
             if dep.tlm_value is None:
                 continue
-            if dep.output is self.forward_cache.func and not self.is_linear:
+            if dep.output is self.forward_cache.func:
                 continue
             relevant_d2Fdm2_forms.append(self.hessian_cache.d2Fdm2_adj_forms[idx][i])
 
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 814748fa3a..660623a2f7 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -17,7 +17,6 @@
         "func",
         "bcs",
         "solver",
-        "is_linear",
         "replaced_deps",
     ]
 )
@@ -149,6 +148,11 @@ def _ad_forward_cache(self):
         Fnew = replace(F, replace_map)
         unew = replace_map[problem.u]
 
+        Jnew = replace(problem.J, replace_map)
+        Jpnew = None
+        if problem.Jp and problem.Jp is not problem.J:
+            Jpnew = replace(problem.Jp, replace_map)
+
         # We also need to "replace" all the bcs in
         # the new NLVS so we can modify those values
         # without affecting user code.
@@ -166,7 +170,7 @@ def _ad_forward_cache(self):
 
         # This NLVS will be used to recompute the solve.
         # TODO: solver_parameters
-        nlvp = NonlinearVariationalProblem(Fnew, unew, bcs=bcs_new)
+        nlvp = NonlinearVariationalProblem(Fnew, unew, J=Jnew, Jp=Jpnew, bcs=bcs_new)
         nlvs = NonlinearVariationalSolver(
             nlvp,
             *self._ad_args_kwargs.forward_args,
@@ -182,7 +186,6 @@ def _ad_forward_cache(self):
             func=self._ad_problem.u,
             bcs=bcs_new,
             solver=nlvs,
-            is_linear=self._ad_problem.is_linear,
             replaced_deps=tuple(replace_map.values()),
         )
 
@@ -190,7 +193,7 @@ def _ad_forward_cache(self):
     @no_annotations
     def _ad_tangent_cache(self):
         from firedrake import (
-            Function, Cofunction, derivative, TrialFunction,
+            Function, Cofunction, derivative,
             LinearVariationalProblem, LinearVariationalSolver)
 
         # If we build the TLM form from the cached
@@ -207,7 +210,7 @@ def _ad_tangent_cache(self):
         # Then for the _partial_ derivatives:
         # (dF/du)*(du/dm) + dF/dm = 0 so we calculate:
         # (dF/du)*(du/dm) = -dF/dm
-        dFdu = derivative(F, u, TrialFunction(V))
+        dFdu = expand_derivatives(derivative(F, u))
         dFdm = Cofunction(V.dual())
         dudm = Function(V)
 
@@ -231,8 +234,6 @@ def _ad_tangent_cache(self):
             replaced_tlms.append(mtlm)
 
             dFdm = derivative(-F, m, mtlm)
-            # TODO: Do we need expand_derivatives here? If so, why?
-            dFdm = expand_derivatives(dFdm)
             dFdm_tlm_forms.append(dFdm)
 
         tlm_val = Function(V)
@@ -336,34 +337,33 @@ def _ad_adjoint_cache(self):
     @no_annotations
     def _ad_hessian_cache(self):
         from firedrake import (
-            Function, TestFunction)
+            Function, TrialFunction)
 
         nlvp = self._ad_forward_cache.solver._problem
         F = nlvp.F
         u = nlvp.u
         V = u.function_space()
 
-        # 1. Forms to calculate rhs of Hessian solve
+        # Solution of the adjoint equation, requires evaluate_adj to
+        # have been called
+        adj_sol = Function(V)
+        # Solution of the TLM equation, requires evaluate_tlm to
+        # have been called (which is handled by the driver)
+        tlm_output = Function(V)
+        # Solution of the second-order adjoint equation,
+        # this is computed during prepare_evaluate_hessian
+        adj2_sol = Function(V)
 
-        # Calculate d^2F/du^2 * du/dm * dm
+        # 1. Forms to calculate rhs of Hessian solve
+        # Calculate d^2F*/du^2 * du/dm * dm
         # where dm is direction for tlm action so du/dm * dm is tlm output
         dFdu = self._ad_tangent_cache.dFdu
-        tlm_output = Function(V)
-        d2Fdu2 = derivative(dFdu, u, tlm_output)
-        # print()
-        # print(f"{dFdu = }")
-        # print()
-        # print(f"{d2Fdu2 = }")
-        # print()
-        d2Fdu2 = expand_derivatives(d2Fdu2)
+        dFdu_adj = action(adjoint(dFdu), adj_sol)
 
-        adj_sol = Function(V)
-
-        # Contribution from tlm_output
-        if len(d2Fdu2.integrals()) > 0:
-            d2Fdu2_form = action(adjoint(d2Fdu2), adj_sol)
-        else:
-            d2Fdu2_form = d2Fdu2
+        d2Fdu2 = derivative(dFdu_adj, u, tlm_output)
+        # expand_derivatives will simplify down to an empty
+        # form if required
+        d2Fdu2_form = expand_derivatives(d2Fdu2)
 
         # Contributions from each tlm_input
         dFdu_adj = action(self._ad_adjoint_cache.dFdu_adj, adj_sol)
@@ -376,27 +376,21 @@ def _ad_hessian_cache(self):
             d2Fdmdu_forms.append(d2Fdmdu)
 
         # 2. Forms to calculate contribution from each control
-        adj2_sol = Function(V)
-
-        Fadj = action(F, adj_sol)
-        Fadj2 = action(F, adj2_sol)
-
         dFdm_adj2_forms = []
         d2Fdm2_adj_forms = []
         d2Fdudm_forms = []
         for m in self._ad_forward_cache.replaced_deps:
-            dm = TestFunction(m.function_space())
-            dFdm_adj2 = expand_derivatives(
-                derivative(Fadj2, m, dm))
+            dm = TrialFunction(m.function_space())
+            dFdm = derivative(F, m, dm)
 
+            dFdm_adj2 = action(adjoint(dFdm), adj2_sol)
             dFdm_adj2_forms.append(dFdm_adj2)
 
-            dFdm_adj = derivative(Fadj, m, dm)
-
-            d2Fdudm = expand_derivatives(
-                derivative(dFdm_adj, u, tlm_output))
-
-            d2Fdudm_forms.append(d2Fdudm)
+            # we need to expand derivatives before taking
+            # the second derivative
+            dFdm_adj = expand_derivatives(action(adjoint(dFdm), adj_sol))
+            d2Fdudm = derivative(dFdm_adj, u, tlm_output)
+            d2Fdudm_forms.append(expand_derivatives(d2Fdudm))
 
             d2Fdm2_adj_forms_k = []
             for m2, dm2 in zip(self._ad_forward_cache.replaced_deps,

From b2f7998caf50f95d13dcd32dfdb8df1ffbd08422 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Wed, 8 Jul 2026 12:15:08 +1000
Subject: [PATCH 6/8] Remove explicit annotation on solve()

We remove the explicit adjoint annotation on the bare solve() call,
relying on the path through the NLVS as the single source of
annotation. This simplifies and unifies the code paths, and will
ideally allow us to remove the specialised solve blocks.

Add CachedSolverBlock to get_solve_blocks check

I don't think this affects anything, but should be there as a
convenience function (referenced in the "extract adjoint solutions"
notebook).

Remove SolveLinearSystemBlock

This is redundant now that it gets routed through the usual
cached solver path.
---
 firedrake/adjoint_utils/__init__.py           |  2 +-
 firedrake/adjoint_utils/blocks/__init__.py    |  2 +-
 firedrake/adjoint_utils/blocks/solving.py     | 19 -----
 firedrake/adjoint_utils/solving.py            | 72 +------------------
 firedrake/adjoint_utils/variational_solver.py | 11 ++-
 firedrake/solving.py                          |  2 -
 tests/firedrake/adjoint/test_solving.py       | 31 ++++++++
 7 files changed, 46 insertions(+), 93 deletions(-)

diff --git a/firedrake/adjoint_utils/__init__.py b/firedrake/adjoint_utils/__init__.py
index e1a31c8f1b..e9096735a4 100644
--- a/firedrake/adjoint_utils/__init__.py
+++ b/firedrake/adjoint_utils/__init__.py
@@ -13,7 +13,7 @@
 from firedrake.adjoint_utils.variational_solver import (  # noqa F401
     NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin
 )
-from firedrake.adjoint_utils.solving import annotate_solve, get_solve_blocks  # noqa F401
+from firedrake.adjoint_utils.solving import get_solve_blocks  # noqa F401
 from firedrake.adjoint_utils.mesh import MeshGeometryMixin  # noqa F401
 from firedrake.adjoint_utils.checkpointing import (  # noqa F401
     enable_disk_checkpointing, disk_checkpointing,
diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py
index 9e4c294b39..fa10ff29fa 100644
--- a/firedrake/adjoint_utils/blocks/__init__.py
+++ b/firedrake/adjoint_utils/blocks/__init__.py
@@ -1,6 +1,6 @@
 from firedrake.adjoint_utils.blocks.assembly import AssembleBlock  # noqa F401
 from firedrake.adjoint_utils.blocks.solving import (  # noqa F401
-    CachedSolverBlock, GenericSolveBlock, SolveLinearSystemBlock,
+    CachedSolverBlock, GenericSolveBlock,
     ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock,
     NonlinearVariationalSolveBlock
 )
diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index 5739c6041a..aaa049b616 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -864,25 +864,6 @@ def solve_init_params(self, args, kwargs, varform):
             self.assemble_kwargs["appctx"] = kwargs["appctx"]
 
 
-class SolveLinearSystemBlock(GenericSolveBlock):
-    def __init__(self, A, u, b, *args, **kwargs):
-        lhs = A.form
-        func = u.function
-        rhs = b.form
-        bcs = A.bcs if hasattr(A, "bcs") else []
-        super().__init__(lhs, rhs, func, bcs, *args, **kwargs)
-
-        # Set up parameters initialization
-        self.ident_zeros_tol = \
-            A.ident_zeros_tol if hasattr(A, "ident_zeros_tol") else None
-        self.assemble_system = \
-            A.assemble_system if hasattr(A, "assemble_system") else False
-
-    def _init_solver_parameters(self, args, kwargs):
-        super()._init_solver_parameters(args, kwargs)
-        solve_init_params(self, args, kwargs, varform=False)
-
-
 class SolveVarFormBlock(GenericSolveBlock):
     def __init__(self, equation, func, bcs=[], *args, **kwargs):
         lhs = equation.lhs
diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py
index e9556e06ea..c7301288f4 100644
--- a/firedrake/adjoint_utils/solving.py
+++ b/firedrake/adjoint_utils/solving.py
@@ -1,71 +1,5 @@
-from functools import wraps
-from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape
-
-from firedrake.adjoint_utils.blocks import SolveVarFormBlock, SolveLinearSystemBlock, GenericSolveBlock, ProjectBlock
-import ufl
-
-
-def annotate_solve(solve):
-    """This solve routine wraps the Firedrake :func:`.solve` call. Its purpose is to annotate the model,
-    recording what solves occur and what forms are involved, so that the adjoint and tangent linear models may be
-    constructed automatically by pyadjoint.
-
-    To disable the annotation, just pass ``annotate=False`` to this routine, and it acts exactly like the
-    Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
-    for the purposes of the adjoint computation (such as projecting fields to other function spaces
-    for the purposes of visualisation).
-
-    The overloaded solve takes optional callback functions to extract adjoint solutions.
-    All of the callback functions follow the same signature, taking a single argument of type Function.
-
-    Keyword Args:
-        adj_cb (:obj:`firedrake.function`, optional):
-            callback function supplying the adjoint solution in the interior. The boundary values are zero.
-        adj_bdy_cb (:obj:`firedrake.function`, optional):
-            callback function supplying the adjoint solution on the boundary.
-            The interior values are not guaranteed to be zero.
-        adj2_cb (:obj:`firedrake.function`, optional):
-            callback function supplying the second-order adjoint solution in the interior.
-            The boundary values are zero.
-        adj2_bdy_cb (:obj:`firedrake.function`, optional):
-            callback function supplying the second-order adjoint solution on
-            the boundary. The interior values are not guaranteed to be zero.
-        ad_block_tag (:obj:`string`, optional):
-            tag used to label the resulting block on the Pyadjoint tape. This
-            is useful for identifying which block is associated with which equation in the forward model.
-
-    """
-
-    @wraps(solve)
-    def wrapper(*args, **kwargs):
-
-        ad_block_tag = kwargs.pop("ad_block_tag", None)
-        annotate = annotate_tape(kwargs)
-
-        if annotate:
-            tape = get_working_tape()
-            solve_block_type = SolveVarFormBlock
-            if not isinstance(args[0], ufl.equation.Equation):
-                solve_block_type = SolveLinearSystemBlock
-
-            sb_kwargs = solve_block_type.pop_kwargs(kwargs)
-            sb_kwargs.update(kwargs)
-            block = solve_block_type(*args, ad_block_tag=ad_block_tag, **sb_kwargs)
-            tape.add_block(block)
-
-        with stop_annotating():
-            output = solve(*args, **kwargs)
-
-        if annotate:
-            if hasattr(args[1], "create_block_variable"):
-                block_variable = args[1].create_block_variable()
-            else:
-                block_variable = args[1].function.create_block_variable()
-            block.add_output(block_variable)
-
-        return output
-
-    return wrapper
+from pyadjoint.tape import get_working_tape
+from firedrake.adjoint_utils.blocks import CachedSolverBlock, GenericSolveBlock, ProjectBlock
 
 
 def get_solve_blocks():
@@ -77,6 +11,6 @@ def get_solve_blocks():
     return [
         block
         for block in get_working_tape().get_blocks()
-        if issubclass(type(block), GenericSolveBlock)
+        if issubclass(type(block), (CachedSolverBlock, GenericSolveBlock))
         and not issubclass(type(block), ProjectBlock)
     ]
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 660623a2f7..1466a6d83c 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -128,6 +128,7 @@ def wrapper(self, problem, *args, **kwargs):
     def _ad_forward_cache(self):
         from firedrake import (
             DirichletBC,
+            MatrixBase,
             NonlinearVariationalProblem,
             NonlinearVariationalSolver)
 
@@ -168,9 +169,17 @@ def _ad_forward_cache(self):
             for bc in bcs
         ]
 
+        # for a pre-assembled Jacobian (i.e. solving an
+        # assembled equation), we can't pass BCs into the
+        # solver, but we need to cache them for the TLM
+        # and adjoint solvers
+        bcs_fwd = []
+        if not isinstance(Jnew, MatrixBase):
+            bcs_fwd = bcs_new
+
         # This NLVS will be used to recompute the solve.
         # TODO: solver_parameters
-        nlvp = NonlinearVariationalProblem(Fnew, unew, J=Jnew, Jp=Jpnew, bcs=bcs_new)
+        nlvp = NonlinearVariationalProblem(Fnew, unew, J=Jnew, Jp=Jpnew, bcs=bcs_fwd)
         nlvs = NonlinearVariationalSolver(
             nlvp,
             *self._ad_args_kwargs.forward_args,
diff --git a/firedrake/solving.py b/firedrake/solving.py
index 414870eba9..c4b80bd4f4 100644
--- a/firedrake/solving.py
+++ b/firedrake/solving.py
@@ -24,13 +24,11 @@
 import firedrake.linear_solver as ls
 import firedrake.variational_solver as vs
 from firedrake.function import Function
-from firedrake.adjoint_utils import annotate_solve
 from firedrake.petsc import PETSc
 from firedrake.utils import ScalarType
 
 
 @PETSc.Log.EventDecorator()
-@annotate_solve
 def solve(*args, **kwargs):
     r"""Solve linear system Ax = b or variational problem a == L or F == 0.
 
diff --git a/tests/firedrake/adjoint/test_solving.py b/tests/firedrake/adjoint/test_solving.py
index 05b114ad0e..d3bca15d39 100644
--- a/tests/firedrake/adjoint/test_solving.py
+++ b/tests/firedrake/adjoint/test_solving.py
@@ -42,6 +42,37 @@ def J(f):
     _test_adjoint(J, f, rg)
 
 
+@pytest.mark.skipcomplex
+def test_assembled_problem(rg):
+    assert len(get_working_tape()._blocks) == 0
+    mesh = IntervalMesh(10, 0, 1)
+    V = FunctionSpace(mesh, "Lagrange", 1)
+    R = FunctionSpace(mesh, "R", 0)
+    f = Function(V).assign(1.)
+
+    u = TrialFunction(V)
+    u_ = Function(V)
+    v = TestFunction(V)
+    bc = DirichletBC(V, Function(R, val=1), "on_boundary")
+
+    def J(f):
+        a = inner(grad(u), grad(v))*dx
+        L = f*v*dx
+
+        A = assemble(a, bcs=bc)
+        b = assemble(L)
+
+        solve(A, u_, b)
+        return assemble(u_**2*dx)
+
+    # this tests a bug in recompute() if bcs= keyword argument is provided in solve
+    J0 = J(f)
+    rf = ReducedFunctional(J0, Control(f))
+    assert_approx_equal(rf(f), J0)
+    assert rf.tape.recompute_count == 1
+    _test_adjoint(J, f, rg)
+
+
 @pytest.mark.skipcomplex
 def test_singular_linear_problem(rg):
     """This tests whether nullspace and solver_parameters are passed on in adjoint solves"""

From 4e1e9d6b938db62f24086cbffadece100a779df6 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Thu, 18 Jun 2026 12:52:44 +1000
Subject: [PATCH 7/8] Handle MeshGeometry dependencies on cached solver blocks

Push subtraction out of derivative

For whatever reason, derivative(-F, m, dm) has very different behaviour
when passed through the adjoint -> action -> expand_derivatives chain
compared to the non-negated variant. We just push subtraction to the
dFdm_adj and dFdm_adj2 forms instead.
---
 firedrake/adjoint_utils/blocks/solving.py     | 23 +++--
 firedrake/adjoint_utils/variational_solver.py | 83 ++++++++++++++++---
 2 files changed, 90 insertions(+), 16 deletions(-)

diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py
index aaa049b616..a2d8582aa6 100644
--- a/firedrake/adjoint_utils/blocks/solving.py
+++ b/firedrake/adjoint_utils/blocks/solving.py
@@ -40,8 +40,9 @@ class SolverType(Enum):
 
 
 class CachedSolverBlock(Block):
-    def __init__(self, forward_cache, tangent_cache, adjoint_cache, hessian_cache,
-                 ad_block_tag=None):
+    def __init__(
+        self, forward_cache, tangent_cache, adjoint_cache, hessian_cache, ad_block_tag=None
+    ):
         super().__init__(ad_block_tag=ad_block_tag)
 
         self.forward_cache = forward_cache
@@ -58,6 +59,12 @@ def _coefficient_dependencies(self, dependencies=None):
         dependencies = dependencies or self.get_dependencies()
         return dependencies[:len(self.forward_cache.replaced_deps)]
 
+    def _mesh_dependencies(self, dependencies=None):
+        dependencies = dependencies or self.get_dependencies()
+        len_replaced = len(self.forward_cache.replaced_deps)
+        len_meshes = len(self.forward_cache.meshes)
+        return dependencies[len_replaced:len_replaced + len_meshes]
+
     def _bc_dependencies(self, dependencies=None):
         dependencies = dependencies or self.get_dependencies()
         if len(self.forward_cache.bcs) > 0:
@@ -98,6 +105,12 @@ def update_tlm_dependencies(self):
                 continue
             replaced_dep.assign(dep.tlm_value)
 
+        for replaced_dep, dep in zip(self.tangent_cache.mesh_tlms,
+                                     self._mesh_dependencies()):
+            if dep.tlm_value is None:
+                continue
+            replaced_dep.assign(dep.tlm_value)
+
         for replaced_dep, dep in zip(self.forward_cache.bcs, self._bc_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 bc_val = 0
@@ -243,7 +256,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
 
         # tlm_input contribution
         for d2Fdmdu, dep in zip(self.hessian_cache.d2Fdmdu_forms,
-                                self._coefficient_dependencies()):
+                                self._coefficient_dependencies() + self._mesh_dependencies()):
             if dep.tlm_value is None:  # This dependency doesn't depend on the controls
                 continue
             if dep.output is self.forward_cache.func:  # Can't compute dependence on initial guess
@@ -279,7 +292,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
 
         relevant_d2Fdm2_forms = []
         for i, dep in relevant_dependencies:
-            if i >= len(self._coefficient_dependencies()):
+            if i >= len(self._coefficient_dependencies() + self._mesh_dependencies()):
                 continue
             if dep.tlm_value is None:
                 continue
@@ -293,7 +306,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v
                      self.hessian_cache.dFdm_adj2_forms[idx],
                      *relevant_d2Fdm2_forms):
             if not form.empty():
-                hessian_output += firedrake.assemble(-form)
+                hessian_output += firedrake.assemble(form)
 
         return hessian_output
 
diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py
index 1466a6d83c..3c8f9ce49a 100644
--- a/firedrake/adjoint_utils/variational_solver.py
+++ b/firedrake/adjoint_utils/variational_solver.py
@@ -7,6 +7,7 @@
 from firedrake.ufl_expr import derivative, adjoint, action
 from ufl import replace, Action
 from ufl.algorithms import expand_derivatives
+from ufl.domain import extract_domains
 from types import SimpleNamespace
 from collections import namedtuple
 
@@ -16,6 +17,7 @@
     field_names=[
         "func",
         "bcs",
+        "meshes",
         "solver",
         "replaced_deps",
     ]
@@ -41,6 +43,7 @@
         "dFdm_forms",
         "solver",
         "replaced_tlms",
+        "mesh_tlms",
         "dFdu",
     ]
 )
@@ -177,6 +180,15 @@ def _ad_forward_cache(self):
         if not isinstance(Jnew, MatrixBase):
             bcs_fwd = bcs_new
 
+        # get all the unique meshes for domains on the form
+        meshes = set()
+        try:
+            for mesh in extract_domains(F):
+                meshes.add(mesh)
+        except AttributeError:
+            pass
+        meshes = list(meshes)
+
         # This NLVS will be used to recompute the solve.
         # TODO: solver_parameters
         nlvp = NonlinearVariationalProblem(Fnew, unew, J=Jnew, Jp=Jpnew, bcs=bcs_fwd)
@@ -189,11 +201,12 @@ def _ad_forward_cache(self):
         # dependencies to all solve blocks.
         # The block need handles to the newly created
         # objects to update their values when recomputing.
-        self._ad_dependencies_to_add = (*replace_map.keys(), *bcs)
+        self._ad_dependencies_to_add = (*replace_map.keys(), *meshes, *bcs)
 
         return ForwardSolveRecomputeCache(
             func=self._ad_problem.u,
             bcs=bcs_new,
+            meshes=meshes,
             solver=nlvs,
             replaced_deps=tuple(replace_map.values()),
         )
@@ -203,7 +216,9 @@ def _ad_forward_cache(self):
     def _ad_tangent_cache(self):
         from firedrake import (
             Function, Cofunction, derivative,
-            LinearVariationalProblem, LinearVariationalSolver)
+            LinearVariationalProblem, LinearVariationalSolver,
+            SpatialCoordinate,
+        )
 
         # If we build the TLM form from the cached
         # forward solve form then we can update exactly
@@ -245,6 +260,15 @@ def _ad_tangent_cache(self):
             dFdm = derivative(-F, m, mtlm)
             dFdm_tlm_forms.append(dFdm)
 
+        mesh_tlms = []
+        for m in self._ad_forward_cache.meshes:
+            mtlm = Function(m._ad_function_space())
+            mesh_tlms.append(mtlm)
+
+            X = SpatialCoordinate(m)
+            dFdm = derivative(-F, X, mtlm)
+            dFdm_tlm_forms.append(dFdm)
+
         tlm_val = Function(V)
 
         return TangentSolveRecomputeCache(
@@ -253,6 +277,7 @@ def _ad_tangent_cache(self):
             dFdm_forms=dFdm_tlm_forms,
             solver=lvs,
             replaced_tlms=replaced_tlms,
+            mesh_tlms=mesh_tlms,
             dFdu=dFdu,
         )
 
@@ -261,7 +286,9 @@ def _ad_tangent_cache(self):
     def _ad_adjoint_cache(self):
         from firedrake import (
             Function, Cofunction, TrialFunction, Argument,
-            LinearVariationalProblem, LinearVariationalSolver)
+            LinearVariationalProblem, LinearVariationalSolver,
+            SpatialCoordinate, TestFunction,
+        )
 
         # If we build the adjoint form from the cached
         # forward solve form then we can update exactly
@@ -327,6 +354,13 @@ def _ad_adjoint_cache(self):
 
             dFdm_adj_forms.append(dFdm)
 
+        for m in self._ad_forward_cache.meshes:
+            X = SpatialCoordinate(m)
+            # we can't take the CoordinateDerivative of an Action, so we have
+            # to invert this form compared to the expression above
+            dFdm = derivative(action(-F, adj_sol), X, TestFunction(m._ad_function_space()))
+            dFdm_adj_forms.append(dFdm)
+
         # To calculate the adjoint component of each DirichletBC
         # we'll need the residual of the adjoint equation without
         # any DirichletBC using the solution calculated with
@@ -346,7 +380,9 @@ def _ad_adjoint_cache(self):
     @no_annotations
     def _ad_hessian_cache(self):
         from firedrake import (
-            Function, TrialFunction)
+            Function, TrialFunction, TestFunction,
+            SpatialCoordinate, MeshGeometry,
+        )
 
         nlvp = self._ad_forward_cache.solver._problem
         F = nlvp.F
@@ -384,20 +420,39 @@ def _ad_hessian_cache(self):
 
             d2Fdmdu_forms.append(d2Fdmdu)
 
+        for m, dm in zip(self._ad_forward_cache.meshes,
+                         self._ad_tangent_cache.mesh_tlms):
+            X = SpatialCoordinate(m)
+            d2Fdmdu = expand_derivatives(
+                derivative(dFdu_adj, X, dm)
+            )
+
+            d2Fdmdu_forms.append(d2Fdmdu)
+
         # 2. Forms to calculate contribution from each control
         dFdm_adj2_forms = []
         d2Fdm2_adj_forms = []
         d2Fdudm_forms = []
-        for m in self._ad_forward_cache.replaced_deps:
-            dm = TrialFunction(m.function_space())
-            dFdm = derivative(F, m, dm)
+        for m in [*self._ad_forward_cache.replaced_deps, *self._ad_forward_cache.meshes]:
+            if isinstance(m, MeshGeometry):
+                X = SpatialCoordinate(m)
+                dm = TestFunction(m._ad_function_space())
+
+                F_adj = action(-F, adj_sol)
+                dFdm_adj = derivative(F_adj, X, dm)
+
+                F_adj2 = action(-F, adj2_sol)
+                dFdm_adj2 = derivative(F_adj2, X, dm)
+            else:
+                dm = TrialFunction(m.function_space())
+                # XXX should we try inverting this back to the way it was before?
+                dFdm = derivative(F, m, dm)
+
+                dFdm_adj = -expand_derivatives(action(adjoint(dFdm), adj_sol))
+                dFdm_adj2 = -action(adjoint(dFdm), adj2_sol)
 
-            dFdm_adj2 = action(adjoint(dFdm), adj2_sol)
             dFdm_adj2_forms.append(dFdm_adj2)
 
-            # we need to expand derivatives before taking
-            # the second derivative
-            dFdm_adj = expand_derivatives(action(adjoint(dFdm), adj_sol))
             d2Fdudm = derivative(dFdm_adj, u, tlm_output)
             d2Fdudm_forms.append(expand_derivatives(d2Fdudm))
 
@@ -408,6 +463,12 @@ def _ad_hessian_cache(self):
                     derivative(dFdm_adj, m2, dm2))
                 d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
 
+            for m2, dm2 in zip(self._ad_forward_cache.meshes,
+                               self._ad_tangent_cache.mesh_tlms):
+                X = SpatialCoordinate(m2)
+                d2Fdm2_adj = expand_derivatives(derivative(dFdm_adj, X, dm2))
+                d2Fdm2_adj_forms_k.append(d2Fdm2_adj)
+
             d2Fdm2_adj_forms.append(d2Fdm2_adj_forms_k)
 
         return HessianSolveRecomputeCache(

From 4c199beed3485b1f27831bd1a2b0542324c6aa34 Mon Sep 17 00:00:00 2001
From: Angus Gibson <angus@agibson.me>
Date: Mon, 13 Jul 2026 14:06:44 +1000
Subject: [PATCH 8/8] Fix PDE-constrained optimisation notebook

---
 docs/notebooks/06-pde-constrained-optimisation.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/docs/notebooks/06-pde-constrained-optimisation.py b/docs/notebooks/06-pde-constrained-optimisation.py
index 2e38da54e6..e57917fb7d 100644
--- a/docs/notebooks/06-pde-constrained-optimisation.py
+++ b/docs/notebooks/06-pde-constrained-optimisation.py
@@ -203,12 +203,13 @@
 print("Jhat(g) = %.8g\nJhat(g_opt) = %.8g" % (Jhat(g), Jhat(g_opt)))
 
 # %% [markdown]
-# To see the optimised flow field, we solve the same problem again, only with the new (optimised) value for the boundary data on $\Gamma_\text{circ}$.  This time we're not interested in annotating the solve, so we tell `firedrake-adjoint` to ignore it by passing `annotate=False`.
+# To see the optimised flow field, we solve the same problem again, only with the new (optimised) value for the boundary data on $\Gamma_\text{circ}$.  This time we're not interested in annotating the solve, so we tell `firedrake-adjoint` to ignore it by running it in a `stop_annotating()` context.
 
 # %%
 g.assign(g_opt)
 w_opt = Function(W)
-solve(a == L, w_opt, bcs=bcs, annotate=False)
+with stop_annotating():
+    solve(a == L, w_opt, bcs=bcs)
 
 # %% tags=["nbval-ignore-output"]
 u_opt, p_opt = w_opt.subfunctions