diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 720f67c..5e6d69f 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -27,5 +27,6 @@ repos:
     rev: v2.3.0
     hooks:
       - id: mypy
+        language: system
         files: ^src/|^tests/
         args: ["--config-file", "pyproject.toml"]
diff --git a/demos/time_distributed_control.py b/demos/time_distributed_control.py
index 04ed874..1ad778a 100644
--- a/demos/time_distributed_control.py
+++ b/demos/time_distributed_control.py
@@ -38,9 +38,11 @@ def solve_heat(ctrls):
     v = ufl.TestFunction(V)
 
     f = dolfinx_adjoint.Function(V, name="source")
-    u_0 = dolfinx_adjoint.Function(V, name="solution")
 
-    F = ((u - u_0) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx
+    u_prev = dolfinx_adjoint.Function(V, name="u_prev")
+    uh = dolfinx_adjoint.Function(V, name="solution")
+    dolfinx_adjoint.assign(0.0, uh)
+    F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx
     a, L = ufl.system(F)
     mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
     exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
@@ -48,13 +50,13 @@ def solve_heat(ctrls):
 
     bc = dolfinx.fem.dirichletbc(0.0, exterior_dofs, V)
 
-    j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx)
+    j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((uh - d) ** 2 * ufl.dx)
 
     t_val = float(dt)
     problem = dolfinx_adjoint.LinearProblem(
         a,
         L,
-        u=u_0,
+        u=uh,
         bcs=[bc],
         petsc_options={
             "ksp_type": "preonly",
@@ -81,6 +83,7 @@ def solve_heat(ctrls):
         dolfinx_adjoint.assign(ctrls[t_val], f)
 
         # Update data function
+        dolfinx_adjoint.assign(uh, u_prev)
 
         # Solve PDE
         problem.solve()
@@ -90,12 +93,12 @@ def solve_heat(ctrls):
             weight = 0.5
         else:
             weight = 1
-        j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx)
+        j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((uh - d) ** 2 * ufl.dx)
         # Update time
         t_val += float(dt)
         dolfinx_adjoint.assign(t_val, t)
 
-    return u_0, d, j
+    return uh, d, j
 
 
 u, d, j = solve_heat(ctrls)
@@ -111,16 +114,75 @@ def solve_heat(ctrls):
 J = j + dolfinx_adjoint.assemble_scalar(regularisation)
 m = [pyadjoint.Control(c) for c in ctrls.values()]
 
-
 rf = pyadjoint.ReducedFunctional(J, m)
 
+# Check accuracy of gradient and Hessian using Taylor test
+with pyadjoint.stop_annotating():
+    # Insert this diagnostic section into your demo right after J and m are defined:
+
+    # 1. Generate a non-zero base control point m_pert
+    m_pert = [dolfinx_adjoint.Function(V, name=f"pert_ctrl_{t_val}") for t_val in ctrls.keys()]
+    for c in m_pert:
+        c.x.array[:] = np.random.uniform(0.1, 1.0, size=c.x.array.shape)
+
+    # 2. Define random directions h
+    h = [pyadjoint.Control(dolfinx_adjoint.Function(V)) for _ in m]
+    for hi in h:
+        hi.control.x.array[:] = np.random.uniform(-0.1, 0.1, size=hi.control.x.array.shape)
+
+    print("\n=== 1. Taylor Test at NON-ZERO Control Point ===")
+    min_val_pert = pyadjoint.taylor_test(rf, m_pert, h)
+    print(f"Convergence rate at perturbed point: {min_val_pert:.4f}")
+    rf(m_pert)
+    print("\n=== 2. Second order taylor test at NON-ZERO Control Point ===")
+    dJdm = sum(drfi._ad_dot(hi) for drfi, hi in zip(rf.derivative(), h, strict=True))
+
+    H = rf.hessian([hi.control for hi in h])
+
+    # 2. Iterate and sum the Hessian dot products piecewise
+    dHddu = sum(Hi._ad_dot(hi) for Hi, hi in zip(H, h, strict=True))
+    min_val = pyadjoint.taylor_test(rf, m_pert, h, dJdm=dJdm, Hm=dHddu)
+    print(f"Convergence rate at perturbed point with Hessian: {min_val:.4f}")
+
+    print("\n=== 3. Direct Finite Difference Gradient Verification ===")
+    eps = 1e-6
+
+    # Compute Adjoint Directional Derivative at m_pert
+    rf(m_pert)
+    grad_adj = rf.derivative()
+    adj_dir_deriv = sum(g._ad_dot(hi) for g, hi in zip(grad_adj, h, strict=True))
+
+    # Forward Perturbation J(m + eps*h)
+    m_plus = [dolfinx_adjoint.Function(V) for _ in m_pert]
+    for mp, m_p, hi in zip(m_plus, m_pert, h, strict=True):
+        mp.x.array[:] = m_p.x.array[:] + eps * hi.control.x.array[:]
+    J_plus = float(rf(m_plus))
+
+    # Backward Perturbation J(m - eps*h)
+    m_minus = [dolfinx_adjoint.Function(V) for _ in m_pert]
+    for mm, m_p, hi in zip(m_minus, m_pert, h, strict=True):
+        mm.x.array[:] = m_p.x.array[:] - eps * hi.control.x.array[:]
+    J_minus = float(rf(m_minus))
+
+    # Central Finite Difference
+    fd_dir_deriv = (J_plus - J_minus) / (2 * eps)
+
+    print(f"Adjoint Directional Derivative:    {adj_dir_deriv:.10e}")
+    print(f"Finite Difference Directional Dev: {fd_dir_deriv:.10e}")
+    rel_diff = abs(adj_dir_deriv - fd_dir_deriv) / (abs(fd_dir_deriv) + 1e-15)
+    print(f"Relative Mismatch:                 {rel_diff:.4e}")
+
+    assert rel_diff < 1e-4, f"Adjoint gradient mismatches finite differences! Relative error: {rel_diff}"
+
+# Reset to ensure that we are at the original control point for the optimization
+rf(list(ctrls.values()))
+
 tape = pyadjoint.get_working_tape()
 tape.visualise_dot("test.dot")
 
 opt_ctrls = pyadjoint.minimize(
     rf,
     method="BFGS",
-    # method="Newton-CG",
     options={"maxiter": 100, "disp": True},
 )
 
diff --git a/src/dolfinx_adjoint/blocks/assembly.py b/src/dolfinx_adjoint/blocks/assembly.py
index dc50824..57ccb71 100644
--- a/src/dolfinx_adjoint/blocks/assembly.py
+++ b/src/dolfinx_adjoint/blocks/assembly.py
@@ -188,7 +188,6 @@ def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
             c_rep = block_variable.saved_output
             if coeff in self.form.coefficients():
                 replaced_coeffs[coeff] = c_rep
-
         form = ufl.replace(self.form, replaced_coeffs)
         return form
 
diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index e7ce1ce..06ca6ab 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -13,8 +13,8 @@
 from ufl.algorithms.analysis import traverse_unique_terminals
 
 from ..compat import get_interpolation_points
+from ..types.function import Function, _create_function
 from ..utils import unroll_dofmap
-from ..types.function import _create_function, Function
 
 if typing.TYPE_CHECKING:
     from petsc4py import PETSc
diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py
index 1ff15be..81c544a 100644
--- a/src/dolfinx_adjoint/blocks/solvers.py
+++ b/src/dolfinx_adjoint/blocks/solvers.py
@@ -5,14 +5,140 @@
 from petsc4py import PETSc
 
 import dolfinx.fem.petsc
+import numpy as np
 import pyadjoint
 import ufl
 from dolfinx.fem.function import Function as _Function
 
+from ..compat import compute_form_adjoint
 from ..petsc_utils import LinearAdjointProblem, solve_linear_problem
 from ..types import Function
 from .assembly import _create_vector, _SpecialVector, assemble_compiled_form
 
+type NestedMutableSequence[T] = T | typing.MutableSequence["NestedMutableSequence[T]"]
+type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"]
+
+
+@typing.overload
+def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ...
+@typing.overload
+def assign_mixed_parts[T: NestedSequence[ufl.Form], S: NestedSequence[ufl.Form]](
+    form1: T, form2: S, /
+) -> tuple[T, S]: ...
+def assign_mixed_parts(
+    *form_structs: NestedSequence[ufl.Form],
+) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]:
+    """
+    Recursively assigns mixed-space `part` indices to UFL Test and Trial functions
+    within nested iterables of forms.
+
+    When solving monolithic block systems in FEniCSx, the UFL arguments must be tagged
+    with a `.part()` index corresponding to their block position. For a block matrix
+    (list of lists), the TestFunction corresponds to the row index, and the
+    TrialFunction corresponds to the column index.
+
+    This utility traverses arbitrary nested structures (e.g., a 2D list for the LHS
+    matrix `a` and a 1D list for the RHS vector `L` simultaneously), extracts arguments
+    that lack a part index, builds a unified replacement map, and applies it.
+
+    Args:
+        *form_structs: One or more UFL forms, or nested iterables (lists/tuples) of
+            UFL forms. Passing multiple structures (like `a` and `L`) ensures they
+            share the same replacement map, preventing mismatched compilation.
+
+    Returns:
+        The modified form structures with identical nesting and sequence types, where
+        all unassigned TestFunction and TrialFunction arguments have been mapped.
+        Returns a single structure if one was passed, otherwise returns a tuple.
+    """
+    replace_map = {}
+
+    def _build_map(obj: NestedSequence[ufl.Form], indices: tuple[int, ...]) -> None:
+        """
+        Recursively discover forms, tracking the depth and index of the nesting.
+        `indices` will be `(row,)` for vectors and `(row, col)` for matrices.
+        """
+        if isinstance(obj, ufl.Form):
+            for arg in obj.arguments():
+                # Only map arguments that haven't been assigned a part yet
+                if arg.part() is None and arg not in replace_map:
+                    num = arg.number()
+
+                    # Because num is 0 for TestFunctions and 1 for TrialFunctions,
+                    # it maps perfectly to our nested dimension indices!
+                    # If num < len(indices), we have traversed deep enough to assign it.
+                    if num < len(indices):
+                        replace_map[arg] = ufl.Argument(arg.ufl_function_space(), number=num, part=indices[num])
+
+        elif isinstance(obj, typing.Iterable):
+            for i, item in enumerate(obj):
+                if item is not None:
+                    # Append current topological index to the path and recurse
+                    _build_map(item, indices + (i,))
+
+    def _replace(obj: typing.Any) -> typing.Any:
+        """
+        Recursively rebuild the structure using the populated replace_map,
+        strictly preserving original sequence types (lists vs. tuples).
+        """
+        if isinstance(obj, ufl.Form):
+            return ufl.replace(obj, replace_map)
+        elif isinstance(obj, (list, tuple)):
+            return type(obj)(_replace(item) for item in obj)
+        return obj
+
+    # 1. Build a shared map across all inputs (e.g., ensuring RHS TestFunctions
+    # perfectly match LHS TestFunctions)
+    for struct in form_structs:
+        _build_map(struct, ())
+
+    # 2. If no replacements are needed, exit early to save computation
+    if not replace_map:
+        return form_structs if len(form_structs) > 1 else form_structs[0]
+
+    # 3. Apply the replacements and unpack if necessary
+    replaced = tuple(_replace(struct) for struct in form_structs)
+    return replaced if len(replaced) > 1 else replaced[0]
+
+
+def to_list(data):
+    if isinstance(data, (tuple, list)):
+        return [to_list(item) for item in data]
+    return data
+
+
+def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int) -> typing.Iterable[ufl.Argument]:
+    """Extract all arguments of a given number, sorted by part."""
+    return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part())
+
+
+def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None:
+    """Sum a blocked form into a single form."""
+    # Handle top-level None
+    if form is None:
+        return None
+
+    if isinstance(form, ufl.Form):
+        return form
+
+    elif isinstance(form, typing.Iterable):
+        # Recursively sum items, filtering out Nones
+        valid_forms: list[ufl.Form] = []
+        for fi in form:
+            summed_fi = sum_form(fi)
+            if summed_fi is not None:
+                valid_forms.append(summed_fi)
+
+        # Handle empty case safely
+        if not valid_forms:
+            return None
+
+        # Safely sum without defaulting to integer 0, removing the need for type: ignore
+        return sum(valid_forms[1:], start=valid_forms[0])
+
+    else:
+        raise TypeError(f"Cannot sum form of type {type(form)}")
+
 
 class LinearProblemBlock(pyadjoint.Block):
     """A linear problem that can be used with adjoint methods.
@@ -21,6 +147,7 @@ class LinearProblemBlock(pyadjoint.Block):
     """
 
     _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
+    _tlm_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
     _second_adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
 
     # 2. Overload for the SCALAR case
@@ -86,6 +213,14 @@ def __init__(
         self._adjoint_petsc_options = adjoint_petsc_options
         self._tlm_petsc_options = tlm_petsc_options
         super().__init__(ad_block_tag=ad_block_tag)
+
+        # Collect all arguments in variational forms and replace them with similar
+        # once that is based on a mixed functionspace.
+        if not isinstance(a, ufl.Form):
+            a, L = assign_mixed_parts(a, L)
+            if P is not None:
+                P, _ = assign_mixed_parts(P, L)
+
         self._lhs = a
         self._rhs = L
         self._preconditioner = P
@@ -105,14 +240,38 @@ def __init__(
             self._u = [pyadjoint.create_overloaded_object(ui) for ui in u]
 
         # NOTE: Add mesh and constants as dependencies later on
-        try:
-            for c in self._lhs.coefficients():  # type: ignore
+
+        # To ensure that the solver can be recycled in time dependent loops, the unknown is also added as a dependency
+        # if present in the form.
+        if isinstance(self._u, dolfinx.fem.Function):
+            assert isinstance(self._lhs, ufl.Form)
+            assert isinstance(self._rhs, ufl.Form)
+            if self._u in self._lhs.coefficients() or self._u in self._rhs.coefficients():
+                raise RuntimeError("The unknown function u should not be present in the variational forms a or L.")
+            for c in self._lhs.coefficients():
                 self.add_dependency(c, no_duplicates=True)
-            for c in self._rhs.coefficients():  # type: ignore
+            for c in self._rhs.coefficients():
                 self.add_dependency(c, no_duplicates=True)
-
-        except AttributeError:
-            raise NotImplementedError("Blocked systems not implemented yet.")
+        elif isinstance(self._u, typing.Iterable):
+            for Ai in self._lhs:  # type: ignore
+                for Aij in Ai:
+                    if Aij is not None:
+                        assert isinstance(Aij, ufl.Form)
+                        for c in Aij.coefficients():
+                            if c in self._u:
+                                raise RuntimeError(
+                                    "The unknown function u should not be present in the variational forms a or L."
+                                )
+                            self.add_dependency(c, no_duplicates=True)
+            for part in self._rhs:  # type: ignore
+                for c in part.coefficients():
+                    if c in self._u:
+                        raise RuntimeError(
+                            "The unknown function u should not be present in the variational forms a or L."
+                        )
+                    self.add_dependency(c, no_duplicates=True)
+        else:
+            raise RuntimeError(f"Unknown type for unknown function u={type(self._u)}.")
         self._compiled_lhs = dolfinx.fem.form(
             self._lhs,
             jit_options=jit_options,
@@ -160,16 +319,17 @@ def __init__(
         if isinstance(self._u, dolfinx.fem.Function):
             self._adjoint_solutions = self._u.copy()
             self._second_adjoint_solutions = self._u.copy()
+            self._tlm_solutions = self._u.copy()
         else:
             assert isinstance(self._u, typing.Iterable)
             self._adjoint_solutions = [u.copy() for u in self._u]
             self._second_adjoint_solutions = [u.copy() for u in self._u]
+            self._tlm_solutions = [u.copy() for u in self._u]
 
         self._adjoint_solver = LinearAdjointProblem(
-            self._compute_adjoint(self._lhs),  # type: ignore[arg-type]
+            self._compute_adjoint(sum_form(self._lhs)),  # type: ignore[arg-type]
             self._rhs,  # type: ignore[arg-type]
             bcs=self._bcs,
-            u=self._adjoint_solutions,  # type: ignore[arg-type]
             P=self._preconditioner,  # type: ignore[arg-type]
             form_compiler_options=self._form_compiler_options,
             jit_options=self._jit_options,
@@ -178,6 +338,7 @@ def __init__(
             kind=kind,  # type: ignore[arg-type]
             entity_maps=self._entity_maps,
         )  # type: ignore[misc]
+        self._tlm_solver = None
 
     def _recover_bcs(self):
         bcs = []
@@ -189,56 +350,110 @@ def _recover_bcs(self):
                 bcs.append(c_rep)
         return bcs
 
-    def _create_replace_map(self, form: ufl.Form) -> dict[Function, Function]:
+    def construct_tlm_solver(self):
+        dFdu_form = self._compute_residual_derivative()
+        tlm_solver = LinearAdjointProblem(
+            dFdu_form,  # type: ignore[arg-type]
+            self._rhs,  # type: ignore[arg-type]
+            bcs=self._bcs,
+            u=self._tlm_solutions,  # type: ignore[arg-type]
+            P=self._preconditioner,  # type: ignore[arg-type]
+            form_compiler_options=self._form_compiler_options,
+            jit_options=self._jit_options,
+            petsc_options=self._tlm_petsc_options,
+            petsc_options_prefix=self._petsc_options_prefix,
+            kind=self._kind,  # type: ignore[arg-type]
+            entity_maps=self._entity_maps,
+        )  # type: ignore[misc]
+        return tlm_solver
+
+    def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]:
         """Replace dependencies with latest checkpoint."""
         replace_map = {}
         for block_variable in self.get_dependencies():
             coeff = block_variable.output
-            if coeff in form.coefficients():
-                replace_map[coeff] = block_variable.saved_output
+            if isinstance(form, ufl.Form):
+                if coeff in form.coefficients():
+                    replace_map[coeff] = block_variable.saved_output
+            elif form is None:
+                return {}
+            else:
+                for f in form:
+                    replace_map.update(self._create_replace_map(f))
+        for block_variable in self.get_outputs():
+            coeff = block_variable.output
+            if isinstance(form, ufl.Form):
+                if coeff in form.coefficients():
+                    replace_map[coeff] = block_variable.saved_output
+            elif form is None:
+                return {}
+            else:
+                for f in form:
+                    replace_map.update(self._create_replace_map(f))
         return replace_map
 
-    def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form:
-        """Replace coefficients in the form with saved outputs.
+    def _create_recompute_replace_map(self, inputs: typing.Sequence[typing.Any]) -> dict:
+        """Map original dependency coefficients to active recomputation inputs."""
+        replace_map = {}
+        for block_variable, input_val in zip(self.get_dependencies(), inputs, strict=True):
+            coeff = block_variable.output
+            val = input_val.output if isinstance(input_val, pyadjoint.block_variable.BlockVariable) else input_val
+            replace_map[coeff] = val
+        return replace_map
 
-        Args:
-            form: The UFL form to replace coefficients in.
-        """
-        replace_map = self._create_replace_map(form)
-        return ufl.replace(form, replace_map)
+    @typing.overload
+    def _replace_form_coefficients_recompute(
+        self, form: ufl.Form | NestedSequence[ufl.Form], replace_map: dict
+    ) -> ufl.Form | NestedMutableSequence[ufl.Form | None]: ...
+    @typing.overload
+    def _replace_form_coefficients_recompute(self, form: None, replace_map: dict) -> None: ...
+    def _replace_form_coefficients_recompute(
+        self, form: ufl.Form | NestedSequence[ufl.Form] | None, replace_map: dict
+    ) -> ufl.Form | NestedMutableSequence[ufl.Form | None] | None:
+        """Recursively replace form coefficients for scalar or blocked form structures."""
+        if form is None:
+            return None
+        if isinstance(form, ufl.Form):
+            coeffs_in_form = form.coefficients()
+            sub_map = {k: v for k, v in replace_map.items() if k in coeffs_in_form}
+            return ufl.replace(form, sub_map) if sub_map else form
+        elif isinstance(form, typing.Sequence):
+            return [self._replace_form_coefficients_recompute(f, replace_map) for f in form]
+        else:
+            raise TypeError(f"Cannot replace coefficients in form of type {type(form)}")
 
-    def prepare_recompute_component(self, inputs, relevant_outputs):
+    def prepare_recompute_component(
+        self, inputs: typing.Sequence[typing.Any], relevant_outputs: typing.Sequence[typing.Any]
+    ) -> dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]:
         """Prepare for recomputing the block with different control inputs."""
 
-        # Create initial guess for the KSP solver
-        # Form independnet compilation would make it possible to use the same KSP for all re-evaluations.
-        if isinstance(self._u, Function):
-            initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess")
-        else:
-            initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u]
+        replace_map = self._create_recompute_replace_map(inputs)
 
-        # Replace form coefficients with checkpointed values.
-        # Loop through the dependencies of the lhs and rhs, check if they are in the respective form
-        lhs = self._replace_coefficients_in_form(self._lhs)
-        rhs = self._replace_coefficients_in_form(self._rhs)
+        # 1. Substitute forms using active candidate inputs instead of static tape checkpoints
+        lhs = self._replace_form_coefficients_recompute(self._lhs, replace_map)
+        rhs = self._replace_form_coefficients_recompute(self._rhs, replace_map)
         preconditioner = (
-            self._replace_coefficients_in_form(self._preconditioner) if self._preconditioner is not None else None
+            self._replace_form_coefficients_recompute(self._preconditioner, replace_map)
+            if self._preconditioner is not None
+            else None
         )
+
+        # 2. Recompile UFL forms with candidate inputs
         compiled_lhs = dolfinx.fem.form(
-            lhs,
+            lhs,  # type: ignore[arg-type]
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
         compiled_rhs = dolfinx.fem.form(
-            rhs,
+            rhs,  # type: ignore[arg-type]
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
         compiled_preconditioner = (
             dolfinx.fem.form(
-                preconditioner,
+                preconditioner,  # type: ignore[arg-type]
                 jit_options=self._jit_options,
                 form_compiler_options=self._form_compiler_options,
                 entity_maps=self._entity_maps,
@@ -247,19 +462,39 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
             else None
         )
 
-        # Replace the compiled forms with those with new coefficients.
-        self._forward_solver._a = compiled_lhs
-        self._forward_solver._L = compiled_rhs
-        self._forward_solver._P = compiled_preconditioner
+        # 3. Hot-swap solver forms
+        self._forward_solver._a = compiled_lhs  # type: ignore[assignment]
+        self._forward_solver._L = compiled_rhs  # type: ignore[assignment]
+        self._forward_solver._preconditioner = compiled_preconditioner
         self._forward_solver.bcs = self._bcs
-        self._forward_solver._u = initial_guess
+        self._forward_solver._u = self._u
+
+        # Clear solution vector
+        if isinstance(self._u, dolfinx.fem.Function):
+            self._u.x.array[:] = 0.0
+        else:
+            for ui in self._u:
+                ui.x.array[:] = 0.0
+        # 4. Solve forward state while halting annotation
+        with pyadjoint.stop_annotating():
+            self._forward_solver.solve()
+        return self._u
 
     def recompute_component(
-        self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None
-    ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]:
-        """Recompute the block with the prepared linear problem."""
-        solution = self._forward_solver.solve()
-        return solution
+        self,
+        inputs: typing.Iterable[Function],
+        block_variable: pyadjoint.block_variable.BlockVariable,
+        idx: int,
+        prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function],
+    ) -> dolfinx.fem.Function:
+        """Recompute and return an isolated copy of the solution state."""
+        if isinstance(prepared, dolfinx.fem.Function):
+            assert idx == 0
+            # Return an explicit copy so each tape block gets an isolated state snapshot
+            return prepared.copy()
+        else:
+            assert isinstance(prepared, typing.Iterable)
+            return prepared[idx].copy()
 
     def _should_compute_boundary_adjoint(
         self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]]
@@ -273,226 +508,199 @@ def _should_compute_boundary_adjoint(
         return bdy
 
     @classmethod
-    @typing.overload
-    def _compute_adjoint(
-        cls, form: typing.Sequence[typing.Sequence[ufl.Form]]
-    ) -> typing.Sequence[typing.Sequence[ufl.Form]]: ...
-
-    @classmethod
-    @typing.overload
-    def _compute_adjoint(cls, form: ufl.Form) -> ufl.Form: ...
-
-    @classmethod
-    def _compute_adjoint(
-        cls, form: typing.Union[ufl.Form, typing.Sequence[typing.Sequence[ufl.Form]]]
-    ) -> typing.Union[ufl.Form, typing.Sequence[typing.Sequence[ufl.Form]]]:
+    def _compute_adjoint(cls, form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]] | ufl.Form:
         """
         Compute adjoint of a bilinear form :math:`a(u, v)`, which could be written as a blocked system.
         """
-        if isinstance(form, ufl.Form):
-            return ufl.adjoint(form)
-        else:
-            assert isinstance(form, typing.Iterable)
-            adj_form: list[list[ufl.Form]] = []
-            tmp_form: list[list[ufl.Form]] = []
-            for i, f_i in enumerate(form):
-                tmp_form.append([])
-                adj_form.append([])
-                for j, form_ij in enumerate(f_i):
-                    tmp_form[i].append(ufl.adjoint(form_ij))
-                    adj_form[i].append(ufl.adjoint(form_ij))
-            for i, f_i in enumerate(tmp_form):
-                for j, form_ij in enumerate(f_i):
-                    adj_form[j][i] = form_ij
-            return adj_form
+        return ufl.extract_blocks(compute_form_adjoint(form))
 
-    def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]:
+    def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]:
         """Convert the formulation :math:`a(u, v)=L(v)` into a residual :math:`F(u_b, v) = 0` where
         :math:`u_b` is the solution of the forward problem at the current time and all coefficients are updated.
         """
         # NOTE: Should probably be possible to compile this form once.
         replacement_functions = self.get_outputs()
-        F_form: typing.Union[ufl.Form, list[ufl.Form]] = []
-        if isinstance(self._u, Function):
-            assert len(replacement_functions) == 1, (
-                f"Expected a single output function, got {len(replacement_functions)}"
-            )
-            F_form = ufl.action(self._lhs, replacement_functions[0].saved_output) - self._rhs
-        else:
-            # Blocked formulation (assuming no mixed function-space)
-            assert len(self._u) == len(replacement_functions), (
-                f"Expected {len(self._u)} output functions, got {len(replacement_functions)}"
-            )
-            for i in range(len(self._u)):
-                assert isinstance(F_form, list)
-                res_i = ufl.ZeroBaseForm((self._u[i],))
-                for j in range(len(self._u)):
-                    res_i += ufl.action(self._lhs[i][j], replacement_functions[j].saved_output)  # type: ignore[index]
-                res_i -= self._rhs[i]  # type: ignore[index]
-                F_form.append(res_i)
-        # NOTE: Will fail for blocked systems atm
-        assert isinstance(F_form, ufl.Form)
+        r_funcs = (
+            [r.saved_output for r in replacement_functions]
+            if len(replacement_functions) > 1
+            else replacement_functions[0].saved_output
+        )
+        summed_form = sum_form(self._lhs)
+        F_form = ufl.action(summed_form, r_funcs) - sum_form(self._rhs)
         replacement_map = self._create_replace_map(F_form)
-        if isinstance(self._u, Function):
-            F_form = ufl.replace(F_form, replacement_map)
-        else:
-            assert isinstance(F_form, list)
-            for j in range(len(F_form)):
-                F_form[j] = ufl.replace(F_form[j], replacement_map)
-        return F_form
+        F_form = ufl.replace(F_form, replacement_map)
+        return F_form, replacement_map
 
     def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.Form]]]:
         """Compute the derivative of the residual with respect to the outputs."""
 
-        F_form = self._compute_residual()
-        outputs = [output.saved_output for output in self.get_outputs()]
-        if len(outputs) == 1:
-            assert isinstance(F_form, ufl.Form)
-            dFdu = ufl.derivative(F_form, outputs[0], ufl.TrialFunction(outputs[0].function_space))
-        else:
-            assert isinstance(F_form, list)
-            dFdu = []
-            for i in range(len(outputs)):
-                dFdu.append([])
-                for j in range(len(outputs)):
-                    dFdu[-1].append(ufl.derivative(F_form[i], outputs[j], ufl.TrialFunction(outputs[j].function_space)))
+        res = self._compute_residual()
+        F_form = res[0] if isinstance(res, tuple) else res
+        assert isinstance(F_form, ufl.Form), "Residual form must be a single UFL form."
+
+        outputs = self.get_outputs()
+        # Use r.saved_output directly; no lookup in replacement_map needed!
+        r_funcs = [r.saved_output for r in outputs]
+
+        test_functions = get_sorted_arguments(F_form.arguments(), 0)
+        trial_functions = [
+            ufl.TrialFunction(output.function_space, part=arg.part())
+            for arg, output in zip(test_functions, r_funcs, strict=True)
+        ]
+
+        dFdu = ufl.derivative(F_form, r_funcs, trial_functions)
+
+        if isinstance(self._u, list):
+            return ufl.extract_blocks(dFdu)
         return dFdu
 
     def prepare_evaluate_tlm(
         self, inputs, tlm_inputs, relevant_outputs
     ) -> tuple[typing.Union[list[ufl.Form], ufl.Form], dolfinx.fem.Form]:
-        F_form = self._compute_residual()
-        dFdu_compiled = dolfinx.fem.form(
+
+        F_form, replacement_map = self._compute_residual()
+        if self._tlm_solver is None:
+            self._tlm_solver = self.construct_tlm_solver()
+        # Even if the solver is cached, we need to replace the form, as the output from pyadjoint
+        # is stored in a new function.
+        assert isinstance(self._tlm_solver, LinearAdjointProblem)
+        self._tlm_solver._a = dolfinx.fem.form(
             self._compute_residual_derivative(),
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
-        return F_form, dFdu_compiled  # type: ignore[return-value]
-
-    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function:
-        """Solve the TLM equation for the block variable.
-
-        .. math::
-
-            \frac{\\partial F}{\\partial u} \frac{\\partial u}{\\partial m} = \frac{\\partial F}{\\partial m}
-
-        """
-        # FIXME: Think about blocks later
-        F, dFdu = prepared
-
-        V = self.get_outputs()[idx].output.function_space
+        #  Build RHS (dFdm) for the monolithic system
+        if isinstance(self._u, list):
+            test_funcs = get_sorted_arguments(F_form.arguments(), 0)
+            dFdm = sum([ufl.ZeroBaseForm((test,)) for test in test_funcs])
+        else:
+            test_funcs = [F_form.arguments()[0]]
+            dFdm = ufl.ZeroBaseForm((test_funcs[0],))
 
-        # FIXME: DirichletBC not block variable yet. Required later on. Currently all bcs should be homogenized
-        bcs = []
-        for bc in self._bcs:
-            bcs.append(bc)
-        dFdm = ufl.ZeroBaseForm((ufl.TestFunction(V),))
         for block_variable in self.get_dependencies():
             tlm_value = block_variable.tlm_value
-            # c = block_variable.output
             c_rep = block_variable.saved_output
             if tlm_value is None:
                 continue
+            assert c_rep in replacement_map.values()
+            # Accumulate sensitivities across all block components
+            dFdm += ufl.derivative(-F_form, c_rep, tlm_value)
 
-            dFdm += ufl.derivative(-F, c_rep, tlm_value)
-
-        if isinstance(dFdm, float):
-            v = dFdu.arguments()[0]
-            dFdm = ufl.ZeroBaseForm((v,))
-
+        # Safely wrap zero forms to prevent compilation crashes
         dFdm = ufl.algorithms.expand_derivatives(dFdm)
+        if isinstance(self._u, list):
+            blocks = ufl.extract_blocks(dFdm)
+            if len(blocks) != len(self._u):
+                # Some zero blocks, manually pad with zero forms
+                _dFdm = [ufl.ZeroBaseForm((test,)) for test in test_funcs]
+                for block in blocks:
+                    args = block.arguments()
+                    assert len(args) == 1, "Expected a single test function in the block."
+                    _dFdm[args[0].part()] = block
+                dFdm = _dFdm
+        else:
+            if dFdm == 0 or dFdm.empty():
+                dFdm = ufl.ZeroBaseForm((test_funcs[0],))
+
         dFdm_compiled = dolfinx.fem.form(
             dFdm,
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
-        dudm = dolfinx.fem.Function(V, name="du_dm_tlm_linearblock")
-        # Create and assemble TLM matrix
-        if not hasattr(self, "_A_tlm"):
-            self._A_tlm = dolfinx.fem.petsc.create_matrix(dFdu)
-
-        self._A_tlm.zeroEntries()
-        dolfinx.fem.petsc.assemble_matrix(self._A_tlm, dFdu, bcs=bcs)  # type: ignore[misc,arg-type]
-        self._A_tlm.assemble()
-
-        # Create TLM KSP and attach matrix
-        if not hasattr(self, "_ksp_tlm"):
-            self._ksp_tlm = PETSc.KSP().create(self._A_tlm.getComm())
-        self._ksp_tlm.setOperators(self._A_tlm)
-
-        # Set TLM solver options
-        if self._tlm_petsc_options is not None:
-            prefix = self._petsc_options_prefix + "tlm_"
-            self._ksp_tlm.setOptionsPrefix(prefix)
-            opts = PETSc.Options()
-            opts.prefixPush(prefix)
-            for k, v in self._tlm_petsc_options.items():
-                opts.setValue(k, v)
-            self._ksp_tlm.setFromOptions()
-            opts.prefixPop()
-
-            # For some strange reason delValue doesn't respect prefixes
-            for k, v in self._tlm_petsc_options.items():
-                opts.delValue(f"{prefix}{k}")
-        # Setup preconditioner
-        self._ksp_tlm.setUp()
-
-        b_tlm = dolfinx.fem.create_vector(dolfinx.fem.extract_function_spaces(dFdm_compiled))
-        b_tlm.array[:] = 0.0
-        dolfinx.fem.petsc.assemble_vector(b_tlm.petsc_vec, dFdm_compiled)
 
-        if bcs is not None:
-            # This system should never be "blocked"
-            dolfinx.fem.petsc.apply_lifting(b_tlm.petsc_vec, [dFdu], bcs=[bcs], alpha=0)
-            dolfinx.la.petsc._ghost_update(b_tlm.petsc_vec, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[arg-type]
-            for bc in bcs:
-                bc.set(b_tlm.array, alpha=0)
-        else:
-            dolfinx.la.petsc._ghost_update(b_tlm, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[arg-type]
+        # 3. Assemble RHS Vector utilizing the internal block-allocated vector
+        b_petsc = self._tlm_solver._b
+        with b_petsc.localForm() as b_loc:
+            b_loc.set(0.0)
 
-        # Use the cached solver to skip reallocation and factorization!
-        self._ksp_tlm.solve(b_tlm.petsc_vec, dudm.x.petsc_vec)
-        dudm.x.scatter_forward()
+        dolfinx.fem.petsc.assemble_vector(b_petsc, dFdm_compiled)
+        dolfinx.la.petsc._ghost_update(b_petsc, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)
 
-        # Explicitly free the temporary RHS vector memory
-        b_tlm.petsc_vec.destroy()
+        # 4. Apply Homogeneous Boundary Conditions safely directly to the block vector
+        if self._bcs:
+            try:
+                for bc in self._bcs:
+                    bc.set(b_petsc.array_w, alpha=0.0)
+            except RuntimeError:
+                # FEniCSx throws RuntimeError for flat .set() on a blocked array.
+                # We use the bcs_by_block utility to handle the nested extraction.
+                from dolfinx.fem.bcs import bcs_by_block
 
-        return dudm
+                V_ext = dolfinx.fem.extract_function_spaces(dFdm_compiled)
+                bcs_lift = bcs_by_block(V_ext, self._bcs)
+                dolfinx.fem.petsc.set_bc(b_petsc, bcs_lift, alpha=0.0)
+
+        # 5. Solve the full monolithic TLM system
+        self._tlm_solver.solve()
+
+        return self._tlm_solutions
+
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function:
+        # The system was solved natively in prepare_evaluate_tlm.
+        # Return the corresponding requested sub-function.
+        if isinstance(self._tlm_solutions, list):
+            return self._tlm_solutions[idx]
+        else:
+            assert isinstance(self._tlm_solutions, dolfinx.fem.Function)
+            return self._tlm_solutions
 
     def prepare_evaluate_adj(
         self,
         inputs: typing.Sequence[Function],
         adj_inputs: typing.Sequence[dolfinx.la.Vector],
         relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]],
-    ) -> typing.Union[ufl.Form, typing.Iterable[ufl.Form]]:
+    ) -> tuple[ufl.Form, dict[Function, Function]]:
         """Prepare the block for evaluating the adjoint."""
 
-        # Compute (dF/du[v])* for the linear problem.
-        F_form = self._compute_residual()
-        dFdu = self._compute_residual_derivative()
-        dFdu_adj = self._compute_adjoint(dFdu)
-
         # Extract dJ/du[v] from the adjoint inputs.
-        assert len(adj_inputs) == 1
-        adj_rhs = adj_inputs[0]
-        dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size)
-        dJdu.array[:] = adj_rhs.array[:].copy()
+        if len(adj_inputs) == 1:
+            adj_rhs = adj_inputs[0]
+            dJdu = self._adjoint_solver._b
+            with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc:
+                dJdu_loc.array[:] = adj_rhs_loc.array[:]
+        else:
+            assert len(adj_inputs) == len(self.get_outputs()), (
+                f"Expected {len(self.get_outputs())} adjoint inputs, got {len(adj_inputs)})"
+            )
+            dJdu = self._adjoint_solver._b
+            with dJdu.localForm() as dJdu_loc:
+                dJdu_loc.set(0.0)
+
+            arrs = []
+            for adj_rhs, output in zip(adj_inputs, self.get_outputs()):
+                local_size = output.output.index_map.size_local * output.output.function_space.dofmap.index_map_bs
+                if adj_rhs is None:
+                    arrs.append(np.zeros(local_size, dtype=dolfinx.default_scalar_type))
+                else:
+                    arrs.append(adj_rhs.array[:local_size])
+            dolfinx.la.petsc.assign(arrs, dJdu)
+            dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)  # type: ignore[arg-type]
 
+        # Compute (dF/du[v])* for the linear problem.
+        F_form, replacement_map = self._compute_residual()
+        dFdu = self._compute_residual_derivative()
+        summed = sum_form(dFdu)
+        dFdu_adj = compute_form_adjoint(summed)
+        dFdu_adj = ufl.algorithms.apply_derivatives.apply_derivatives(ufl.algorithms.expand_derivatives(dFdu_adj))
+        assert dFdu_adj.empty() is False, "Adjoint of dF/du[v] is empty. Check if the problem is linear."
         # Solve adjoint problem
         compiled_dFdu = dolfinx.fem.form(
-            dFdu_adj,  # type: ignore[arg-type]
+            ufl.extract_blocks(dFdu_adj),  # type: ignore[arg-type]
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
         self._adjoint_solver._a = compiled_dFdu
-        self._adjoint_solver._b = dJdu.petsc_vec
-        self._adjoint_solver._u = self._adjoint_solutions  # type: ignore[assignment]
         self._adjoint_solver.solve()
-
-        return F_form
+        if isinstance(self._adjoint_solutions, list):
+            for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u):
+                adj_sol.x.array[:] = sol.x.array[:]
+        else:
+            assert isinstance(self._adjoint_solutions, dolfinx.fem.Function)
+            self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:]
+        return F_form, replacement_map
 
     def evaluate_adj_component(
         self,
@@ -500,23 +708,32 @@ def evaluate_adj_component(
         adj_inputs: typing.Iterable[dolfinx.la.Vector],
         block_variable: pyadjoint.block_variable.BlockVariable,
         idx: int,
-        prepared: typing.Union[ufl.Form, typing.Iterable[ufl.Form]],
-    ) -> typing.Union[_SpecialVector, typing.Iterable[_SpecialVector]]:
-        """Evaluate the adjoint component, i.e. :math:`\frac{\\partial Au - b}{\\partial c}`."""
-
-        residual = prepared
+        prepared: tuple[ufl.Form, dict[Function, Function]],
+    ) -> _SpecialVector:
+        """Evaluate the adjoint component, i.e. :math:`\\frac{\\partial F}{\\partial m}`."""
 
+        residual, replacement_map = prepared
         c = block_variable.output
-
         c_rep = block_variable.saved_output
-        if isinstance(c, dolfinx.fem.Function):
-            dc = ufl.TrialFunction(c.function_space)
+        if isinstance(c, Function):
+            # Need some clever construction of the TrialFunction to get a part of the mixed space
+            part = idx if isinstance(self._u, list) else None
+            dc = ufl.TrialFunction(c_rep.function_space, part=part)
         else:
             raise NotImplementedError(f"Unsupported control {type(c)}")
 
-        dFdm = -ufl.derivative(residual, c_rep, dc)
+        # Compute the sensitivity of the residual with respect to the parameter
+        sum_res = sum_form(residual)
+        assert c in replacement_map.keys()
+        assert c_rep == replacement_map[c]
+        dFdm = -ufl.derivative(sum_res, c_rep, dc)
+        if dFdm.empty():
+            # Generate a dummy form to safely extract the correct Vector wrapper type
+            dFdm = dolfinx.fem.form(ufl.ZeroBaseForm((dc,)))  # type: ignore[call-overload]
+
         dFdm_adj = ufl.adjoint(dFdm)
         sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions)
+
         compiled_sensitivity = dolfinx.fem.form(
             sensitivity,
             jit_options=self._jit_options,
@@ -537,15 +754,21 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
 
         # Using the equation Form we derive dF/du, d^2F/du^2 * du/dm * direction.
         dFdu_form = self._compute_residual_derivative()
-        assert len(outputs) == 1, "Hessian computation only implemented for single output blocks."
-        assert len(tlm_output) == 1, "Hessian computation only implemented for single TLM output blocks."
-        d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(dFdu_form, outputs[0].saved_output, tlm_output[0]))
+
+        # For linear forms d2Fdu2 is zero, but we include it for completeness.
+
+        unknowns = [output.saved_output for output in self.get_outputs()]
+        summed_form = sum_form(dFdu_form)
+        d2Fdu2 = ufl.algorithms.expand_derivatives(ufl.derivative(summed_form, unknowns, tlm_output))
 
         # bdy = self._should_compute_boundary_adjoint(relevant_dependencies)
-        assert len(hessian_inputs) == 1, "Hessian computation only implemented for single hessian input blocks."
 
         # Assemble right hand side of second order adjoint equation
+        # Note this term should always be zero for linear problems, but we include it for completeness.
+        if not d2Fdu2.empty():
+            raise RuntimeError(f"This term {d2Fdu2:s} should be zero for linear problems.")
         b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions)
+        dFdu_adj = self._compute_adjoint(sum_form(dFdu_form))
         for bo in self.get_dependencies():
             c = bo.output
             c_rep = bo.saved_output
@@ -555,28 +778,70 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
             if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)):
                 raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.")
             else:
-                dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions)
-                b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input)
+                summed_form = sum_form(dFdu_adj)
+                dFdu_adj_applied = ufl.action(summed_form, self._adjoint_solutions)
+                b_form += ufl.derivative(dFdu_adj_applied, c_rep, tlm_input)
 
-        b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size)
-        b.array[:] = 0.0
-        if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty():
-            compiled_soa_rhs = dolfinx.fem.form(
-                b_form,
-                jit_options=self._jit_options,
-                form_compiler_options=self._form_compiler_options,
-                entity_maps=self._entity_maps,
-            )
-            dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs)
-            b.scatter_reverse(dolfinx.la.InsertMode.add)
-            b.array[:] *= -1
+        if len(outputs) == 1:
+            b = self._adjoint_solver._b
+            with b.localForm() as b_loc:
+                b_loc.set(0.0)
+            form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form)
+            if not form_i.empty():
+                compiled_soa_rhs = dolfinx.fem.form(
+                    form_i,
+                    jit_options=self._jit_options,
+                    form_compiler_options=self._form_compiler_options,
+                    entity_maps=self._entity_maps,
+                )
+                dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs)
+                b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)
+
+                b.scale(-1)
+
+            with b.localForm() as b_loc:
+                b_loc.array[:] += hessian_inputs[0].array[:]
+            b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
+            self._adjoint_solver._b = b
 
-        b.array[:] += hessian_inputs[0].array
-        b.scatter_forward()
+        else:
+            bs = []
+            b_form = ufl.extract_blocks(b_form)
+            for i, hess_input in enumerate(hessian_inputs):
+                if hess_input is not None:
+                    bi = dolfinx.la.vector(hess_input.index_map, hess_input.block_size)
+                else:
+                    out_i = self.get_outputs()[i].saved_output
+                    bi = dolfinx.la.vector(
+                        out_i.function_space.dofmap.index_map, out_i.function_space.dofmap.index_map_bs
+                    )
+                bs.append(bi)
+                bi.array[:] = 0.0
+                form_i = ufl.algorithms.apply_derivatives.apply_derivatives(b_form[i])
+                if not form_i.empty():
+                    compiled_soa_rhs = dolfinx.fem.form(
+                        form_i,
+                        jit_options=self._jit_options,
+                        form_compiler_options=self._form_compiler_options,
+                        entity_maps=self._entity_maps,
+                    )
+                    dolfinx.fem.assemble_vector(bi.array, compiled_soa_rhs)
+                    bi.scatter_reverse(dolfinx.la.InsertMode.add)
+                    bi.scatter_forward()
+                    bi.array[:] *= -1
+
+                if hess_input is not None:
+                    bi.array[:] += hess_input.array
+
+                bi.scatter_forward()
+            b = self._adjoint_solver._b
+            local_arrays = [bi.array[: bi.index_map.size_local * bi.block_size] for bi in bs]
+            dolfinx.la.petsc.assign(local_arrays, b)
+            b.ghostUpdate(PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
 
         # Compile SOA LHS
         dFdu_adj = dolfinx.fem.form(
-            ufl.adjoint(dFdu_form),
+            self._compute_adjoint(sum_form(dFdu_form)),
             jit_options=self._jit_options,
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
@@ -584,9 +849,13 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
 
         # Solve adjoint problem
         self._adjoint_solver._a = dFdu_adj
-        self._adjoint_solver._b = b.petsc_vec
-        self._adjoint_solver._u = self._second_adjoint_solutions
         self._adjoint_solver.solve()
+        if isinstance(self._second_adjoint_solutions, list):
+            for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u):
+                adj_sol.x.array[:] = sol.x.array[:]
+        else:
+            self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:]
+
         return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions
 
     def evaluate_hessian_component(
@@ -601,11 +870,10 @@ def evaluate_hessian_component(
     ):
         c = block_variable.output
 
-        F_form, adj_sol, adj_sol2 = prepared
+        (F_form, replacement_map), adj_sol, adj_sol2 = prepared
 
         outputs = self.get_outputs()
-        assert len(outputs) == 1, "Hessian computation only implemented for single output blocks."
-        tlm_output = outputs[0].tlm_value
+        tlm_output = [output.tlm_value for output in outputs]
 
         c_rep = block_variable.saved_output
 
@@ -627,22 +895,25 @@ def evaluate_hessian_component(
             assert isinstance(c, dolfinx.fem.Function)
             W = c.function_space
 
+        # We are trying to compute (dF/dm)^T lambda_1
+        # and (dF_dm)^T lambda_ 2. However, standard approach of UFL
+        # does not work for MixedFunctionSpaces, as the control space is not
+        # mixed. Therefore, we instead we compute it as dL/dm = d(lambda_i^T F(m))/dm,
+        # which is equivalent.
+        F_summed = sum_form(F_form)
+        L1 = ufl.action(F_summed, adj_sol)
+        L2 = ufl.action(F_summed, adj_sol2)
+
+        # Compute first derivatives (1-forms tested exactly against the single 'dc' object)
         dc = ufl.TestFunction(W)
-        form_adj = ufl.action(F_form, adj_sol)
-        form_adj2 = ufl.action(F_form, adj_sol2)
-        if isinstance(c, dolfinx.mesh.Mesh):
-            raise NotImplementedError("Hessian computation for Mesh control not implemented yet.")
-            # dFdm_adj = ufl.derivative(form_adj, X, dc)
-            # dFdm_adj2 = ufl.derivative(form_adj2, X, dc)
-        else:
-            # Assume Function
-            dFdm_adj = ufl.derivative(form_adj, c_rep, dc)
-            dFdm_adj2 = ufl.derivative(form_adj2, c_rep, dc)
+        assert c_rep in replacement_map.values()
+        dL1dm = ufl.derivative(L1, c_rep, dc)
+        dL2dm = ufl.derivative(L2, c_rep, dc)
 
-        # TODO: Old comment claims this might break on split. Confirm if true or not.
-        d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, outputs[0].saved_output, tlm_output))
+        sa = [output.saved_output for output in outputs]
+        d2Fdudm = ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, sa, tlm_output))
 
-        d2Fdm2 = 0
+        d2Fdm2 = ufl.ZeroBaseForm((dc,))  # Initialize the second derivative form
         # We need to add terms from every other dependency
         # i.e. the terms d^2F/dm_1dm_2
         for _, bv in relevant_dependencies:
@@ -655,19 +926,13 @@ def evaluate_hessian_component(
             if tlm_input is None:
                 continue
 
-            # If problem is non-linear we need to skip the output variable as a control, as we can't differentiate with
-            # respect to the initial guess
-            # if c2 == self._u and not self.linear:
-            #     continue
-
-            # TODO: If tlm_input is a Sum, this crashes in some instances?
             if isinstance(c2_rep, dolfinx.mesh.Mesh):
                 X = ufl.SpatialCoordinate(c2_rep)
-                d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, X, tlm_input))
+                d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, X, tlm_input))
             else:
-                d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dFdm_adj, c2_rep, tlm_input))
+                d2Fdm2 += ufl.algorithms.expand_derivatives(ufl.derivative(dL1dm, c2_rep, tlm_input))
 
-        hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dFdm_adj2 + d2Fdudm)
+        hessian_form = ufl.algorithms.expand_derivatives(d2Fdm2 + dL2dm + d2Fdudm)
 
         compiled_hessian = dolfinx.fem.form(
             hessian_form,
@@ -675,7 +940,11 @@ def evaluate_hessian_component(
             form_compiler_options=self._form_compiler_options,
             entity_maps=self._entity_maps,
         )
-        hessian_output = _create_vector(compiled_hessian, hessian_form.arguments()[0].ufl_function_space())
+
+        test_functions = get_sorted_arguments(hessian_form.arguments(), 0)
+        assert len(test_functions) == 1
+        hessian_output = _create_vector(compiled_hessian, test_functions[0].ufl_function_space())
+
         hessian_output.array[:] = 0.0
         assemble_compiled_form(compiled_hessian, hessian_output)
         hessian_output.array[:] *= -1.0
@@ -690,6 +959,7 @@ class NonlinearProblemBlock(pyadjoint.Block):
 
     _adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
     _second_adjoint_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
+    _tlm_solutions: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
     _rhs: ufl.Form | typing.Sequence[ufl.Form]
 
     @typing.overload
@@ -759,7 +1029,10 @@ def __init__(
         self._u: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
         if isinstance(u, dolfinx.fem.Function):
             self._u = pyadjoint.create_overloaded_object(u)
-            self._rhs = F
+            replace_dict = {u: self._u}
+            self._rhs = ufl.replace(F, replace_dict)
+            self._lhs = ufl.replace(J, replace_dict) if J is not None else None
+            self._preconditioner = ufl.replace(P, replace_dict) if P is not None else None
         else:
             self._u = [pyadjoint.create_overloaded_object(ui) for ui in u]
             assert isinstance(F, typing.Iterable)
@@ -768,12 +1041,20 @@ def __init__(
 
         # NOTE: Add mesh and constants as dependencies later on
         try:
-            for c in self._lhs.coefficients():  # type: ignore
-                self.add_dependency(c, no_duplicates=True)
-            for c in self._rhs.coefficients():  # type: ignore
-                self.add_dependency(c, no_duplicates=True)
+            u_list = self._u if isinstance(self._u, list) else [self._u]
+            if self._lhs is not None:
+                assert isinstance(self._lhs, ufl.Form)
+                for c in self._lhs.coefficients():
+                    if c not in u_list:  # Exclude unknown
+                        self.add_dependency(c, no_duplicates=True)
+            if self._rhs is not None:
+                assert isinstance(self._rhs, ufl.Form)
+                for c in self._rhs.coefficients():
+                    if c not in u_list:  # Exclude unknown
+                        self.add_dependency(c, no_duplicates=True)
         except AttributeError:
             raise NotImplementedError("Blocked systems not implemented yet.")
+
         self._compiled_lhs = dolfinx.fem.form(
             self._lhs,  # type: ignore
             jit_options=jit_options,
@@ -814,10 +1095,12 @@ def __init__(
         if isinstance(self._u, dolfinx.fem.Function):
             self._adjoint_solutions = self._u.copy()  # type: ignore[assignment]
             self._second_adjoint_solutions = self._u.copy()  # type: ignore[assignment]
+            self._tlm_solutions = self._u.copy()  # type: ignore[assignment]
         else:
             assert isinstance(self._u, typing.Iterable)
             self._adjoint_solutions = [u.copy() for u in self._u]
             self._second_adjoint_solutions = [u.copy() for u in self._u]
+            self._tlm_solutions = [u.copy() for u in self._u]
 
         if isinstance(F, ufl.Form):
             dFdu_adj = ufl.adjoint(ufl.derivative(F, u))
@@ -827,7 +1110,6 @@ def __init__(
             dFdu_adj,  # type: ignore[arg-type]
             self._rhs,  # type: ignore[arg-type]
             bcs=self._bcs,
-            u=self._adjoint_solutions,  # type: ignore[arg-type]
             P=self._preconditioner,  # type: ignore[arg-type]
             form_compiler_options=self._form_compiler_options,
             jit_options=self._jit_options,
@@ -868,71 +1150,37 @@ def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form:
     def prepare_recompute_component(self, inputs, relevant_outputs):
         """Prepare for recomputing the block with different control inputs."""
 
-        # Create initial guess for the KSP solver
-        # Form independnet compilation would make it possible to use the same KSP for all re-evaluations.
-        if isinstance(self._u, Function):
-            initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess")
-        else:
-            initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u]
-
-        # Replace values in the DirichletBC if it is dependent on a control
-        # NOTE: Currently assume that BCS are control independent.
-        bcs = self._bcs
-        # for block_variable in self.get_dependencies():
-        #     c = block_variable.output
-        #     c_rep = block_variable.saved_output
-
-        #     if isinstance(c, dolfinx.fem.DirichletBC):
-        #         bcs.append(c_rep)
-
-        # Replace form coefficients with checkpointed values.
-        # Loop through the dependencies of the lhs and rhs, check if they are in the respective form
-        lhs = self._replace_coefficients_in_form(self._lhs)
-        rhs = self._replace_coefficients_in_form(self._rhs)
-        preconditioner = (
-            self._replace_coefficients_in_form(self._preconditioner) if self._preconditioner is not None else None
-        )
-        compiled_lhs = dolfinx.fem.form(
-            lhs,
-            jit_options=self._jit_options,
-            form_compiler_options=self._form_compiler_options,
-            entity_maps=self._entity_maps,
-        )
-        compiled_rhs = dolfinx.fem.form(
-            rhs,
-            jit_options=self._jit_options,
-            form_compiler_options=self._form_compiler_options,
-            entity_maps=self._entity_maps,
-        )
-        compiled_preconditioner = (
-            dolfinx.fem.form(
-                preconditioner,
-                jit_options=self._jit_options,
-                form_compiler_options=self._form_compiler_options,
-                entity_maps=self._entity_maps,
-            )
-            if preconditioner is not None
-            else None
-        )
+        # As opposed to the linear problem, we need to update the coefficients in place,
+        # as the nonlinear problem snes.setContext doesn't reflect in place updates on the solver.
+        for block_variable in self.get_dependencies():
+            coeff = block_variable.output
+            if isinstance(coeff, dolfinx.fem.Function):
+                coeff.x.array[:] = block_variable.saved_output.x.array[:]
+                coeff.x.scatter_forward()
 
-        # Replace the compiled forms with those with new coefficients.
-        self._forward_solver._a = compiled_lhs
-        self._forward_solver._L = compiled_rhs
-        self._forward_solver._P = compiled_preconditioner
-        self._forward_solver.bcs = bcs
-        self._forward_solver._u = initial_guess
+        # Warm-start original unknown objects in place
+        u_list = self._forward_solver._u if isinstance(self._forward_solver._u, list) else [self._forward_solver._u]
+        for idx, out_bv in relevant_outputs:
+            u_list[idx].x.array[:] = out_bv.saved_output.x.array[:]
+            u_list[idx].x.scatter_forward()
+
+        return None
 
     def recompute_component(
         self, inputs: typing.Iterable[Function], block_variable, idx: int, prepared: None
     ) -> typing.Union[dolfinx.fem.Function, typing.Iterable[dolfinx.fem.Function]]:
         """Recompute the block with the prepared linear problem."""
-        solution = self._forward_solver.solve()
-        return solution
+        with pyadjoint.tape.stop_annotating():
+            self._forward_solver.solve()
+
+        if isinstance(self._forward_solver._u, list):
+            return self._forward_solver._u[idx]
+        else:
+            return self._forward_solver._u
 
     def _should_compute_boundary_adjoint(
         self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]]
     ) -> bool:
-        """Determine if the adjoint should be computed with respect to the boundary conditions."""
         bdy = False
         for _, dep in relevant_dependencies:
             if isinstance(dep.output, dolfinx.fem.DirichletBC):
@@ -980,21 +1228,19 @@ def _compute_residual(self) -> typing.Union[ufl.Form, list[ufl.Form]]:
         """
         # NOTE: Should probably be possible to compile this form once.
         replacement_functions = self.get_outputs()
-        F_form: typing.Union[ufl.Form, list[ufl.Form]] = []
+        assert isinstance(self._rhs, (ufl.Form, typing.Sequence))
         assert isinstance(self._rhs, ufl.Form)
         replacement_map = self._create_replace_map(self._rhs)
+
         u_list = self._u if isinstance(self._u, list) else [self._u]
         for u, block in zip(u_list, replacement_functions):
             replacement_map[u] = block.saved_output
 
-        if isinstance(self._u, Function):
+        if isinstance(self._u, dolfinx.fem.Function):
             F_form = ufl.replace(self._rhs, replacement_map)
         else:
-            assert isinstance(F_form, list)
-            assert isinstance(self._rhs, typing.Sequence)
-            assert len(F_form) == len(self._rhs)
-            for j, rhs_j in enumerate(self._rhs):
-                F_form[j] = ufl.replace(rhs_j, replacement_map)
+            assert isinstance(self._rhs, typing.Iterable)
+            F_form = [ufl.replace(rhs_j, replacement_map) for rhs_j in self._rhs]
         return F_form
 
     def _compute_residual_derivative(self) -> typing.Union[ufl.Form, list[list[ufl.Form]]]:
@@ -1035,7 +1281,6 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
             \frac{\\partial F}{\\partial u} \frac{\\partial u}{\\partial m} = \frac{\\partial F}{\\partial m}
 
         """
-        # FIXME: Think about blocks later
         F, dFdu = prepared
 
         V = self.get_outputs()[idx].output.function_space
@@ -1044,14 +1289,13 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         bcs = []
         for bc in self._bcs:
             bcs.append(bc)
+
         dFdm = ufl.ZeroBaseForm((ufl.TestFunction(V),))
         for block_variable in self.get_dependencies():
             tlm_value = block_variable.tlm_value
-            # c = block_variable.output
             c_rep = block_variable.saved_output
             if tlm_value is None:
                 continue
-
             dFdm += ufl.derivative(-F, c_rep, tlm_value)
 
         if isinstance(dFdm, float):
@@ -1098,8 +1342,9 @@ def prepare_evaluate_adj(
         # Extract dJ/du[v] from the adjoint inputs.
         assert len(adj_inputs) == 1
         adj_rhs = adj_inputs[0]
-        dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size)
-        dJdu.array[:] = adj_rhs.array[:].copy()
+        dJdu = self._adjoint_solver._b
+        with dJdu.localForm() as dJdu_loc, adj_rhs.petsc_vec.localForm() as adj_rhs_loc:
+            dJdu_loc.array[:] = adj_rhs_loc.array[:]
 
         # Solve adjoint problem
         compiled_dFdu = dolfinx.fem.form(
@@ -1109,9 +1354,13 @@ def prepare_evaluate_adj(
             entity_maps=self._entity_maps,
         )
         self._adjoint_solver._a = compiled_dFdu
-        self._adjoint_solver._b = dJdu.petsc_vec
-        self._adjoint_solver._u = self._adjoint_solutions  # type: ignore[assignment]
         self._adjoint_solver.solve()
+        if isinstance(self._adjoint_solutions, list):
+            for adj_sol, sol in zip(self._adjoint_solutions, self._adjoint_solver.u):
+                adj_sol.x.array[:] = sol.x.array[:]
+        else:
+            assert isinstance(self._adjoint_solutions, dolfinx.fem.Function)
+            self._adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:]
         return F_form
 
     def evaluate_adj_component(
@@ -1134,6 +1383,11 @@ def evaluate_adj_component(
         else:
             raise NotImplementedError(f"Unsupported control {type(c)}")
         dFdm = -ufl.derivative(residual, c_rep, dc)
+
+        # Safe return for empty sensitivities
+        if dFdm.empty():
+            dFdm = ufl.ZeroBaseForm((dc,))
+
         dFdm_adj = ufl.adjoint(dFdm)
         sensitivity = ufl.action(dFdm_adj, self._adjoint_solutions)
         compiled_sensitivity = dolfinx.fem.form(
@@ -1165,6 +1419,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
 
         # Assemble right hand side of second order adjoint equation
         b_form = d2Fdu2 if d2Fdu2.empty() else ufl.action(ufl.adjoint(d2Fdu2), self._adjoint_solutions)
+        dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions)
         for bo in self.get_dependencies():
             c = bo.output
             c_rep = bo.saved_output
@@ -1174,11 +1429,10 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
             if isinstance(c, (dolfinx.mesh.Mesh, dolfinx.fem.DirichletBC)):
                 raise NotImplementedError(f"Hessian computation for {type(c)} control not implemented yet.")
             else:
-                dFdu_adj = ufl.action(ufl.adjoint(dFdu_form), self._adjoint_solutions)
                 b_form += ufl.derivative(dFdu_adj, c_rep, tlm_input)
-
-        b = dolfinx.la.vector(hessian_inputs[0].index_map, hessian_inputs[0].block_size)
-        b.array[:] = 0.0
+        b = self._adjoint_solver._b
+        with b.localForm() as b_loc:
+            b_loc.set(0.0)
         if not ufl.algorithms.apply_derivatives.apply_derivatives(b_form).empty():
             compiled_soa_rhs = dolfinx.fem.form(
                 b_form,
@@ -1186,11 +1440,11 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
                 form_compiler_options=self._form_compiler_options,
                 entity_maps=self._entity_maps,
             )
-            dolfinx.fem.petsc.assemble_vector(b.petsc_vec, compiled_soa_rhs)
-            b.scatter_reverse(dolfinx.la.InsertMode.add)
-            b.array[:] *= -1
-
-        b.array[:] += hessian_inputs[0].array
+            dolfinx.fem.petsc.assemble_vector(b, compiled_soa_rhs)
+            b.ghostUpdate(PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore [arg-type]
+            b.scale(-1)
+        with b.localForm() as b_loc, hessian_inputs[0].petsc_vec.localForm() as hess_loc:
+            b_loc.array[:] += hess_loc.array[:]
 
         # Compile SOA LHS
         dFdu_adj = dolfinx.fem.form(
@@ -1201,9 +1455,13 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
         )
 
         self._adjoint_solver._a = dFdu_adj
-        self._adjoint_solver._b = b.petsc_vec
         self._adjoint_solver._u = self._second_adjoint_solutions
         self._adjoint_solver.solve()
+        if isinstance(self._second_adjoint_solutions, list):
+            for adj_sol, sol in zip(self._second_adjoint_solutions, self._adjoint_solver.u):
+                adj_sol.x.array[:] = sol.x.array[:]
+        else:
+            self._second_adjoint_solutions.x.array[:] = self._adjoint_solver.u.x.array[:]
         return self._compute_residual(), self._adjoint_solutions, self._second_adjoint_solutions
 
     def evaluate_hessian_component(
diff --git a/src/dolfinx_adjoint/compat.py b/src/dolfinx_adjoint/compat.py
index a040be9..e8b7839 100644
--- a/src/dolfinx_adjoint/compat.py
+++ b/src/dolfinx_adjoint/compat.py
@@ -1,4 +1,9 @@
 import dolfinx
+from ufl.algebra import Conj
+from ufl.algorithms.formsplitter import extract_blocks
+from ufl.algorithms.map_integrands import map_integrands
+from ufl.algorithms.replace import replace
+from ufl.argument import Argument
 
 try:
     from ufl.algorithms.extract_linear_combination import extract_linear_combination
@@ -208,3 +213,89 @@ def get_interpolation_points(V: dolfinx.fem.FunctionSpace):
         return V.element.interpolation_points()  # type: ignore[operator]
     except TypeError:
         return V.element.interpolation_points
+
+
+# Workaround until https://github.com/FEniCS/ufl/pull/508 is in all stable releases we support
+def compute_form_adjoint(
+    form,
+    reordered_arguments: tuple[Argument, Argument] | tuple[tuple[Argument, Argument], ...] | None = None,
+):
+    """Compute the adjoint of a bilinear form.
+
+    This works simply by swapping the number of the two arguments,
+    but keeping their elements and places in the integrand expressions.
+
+    Args:
+        form: A UFL bilinear form.
+        reordered_arguments: Optional explicit arguments to use for the adjoint form.
+            - For standard finite element spaces: A single tuple `(new_u, new_v)`
+              representing the replacement trial and test functions.
+            - For mixed function spaces: A sequence of tuples, with one `(new_u, new_v)`
+              pair for each *subspace*. For example, `((new_u0, new_v0), (new_u1, new_v1))`.
+              The test function mappings are extracted using the block row index `i`,
+              and the trial function mappings using the block column index `j`.
+
+    Returns:
+        The adjoint of the bilinear form.
+    """
+    if form.empty():
+        return form
+
+    arguments = form.arguments()
+
+    # Check if mixed space
+    is_mixed = any(arg.part() is not None for arg in arguments)
+
+    def validate_mapping(old_v: Argument, old_u: Argument, new_v: Argument, new_u: Argument, check_parts=False):
+        """Validate the mapping of old arguments to new arguments."""
+        if new_u.number() >= new_v.number():
+            raise ValueError("Ordering of new arguments is the same as the old arguments!")
+        if new_u.ufl_function_space() != old_u.ufl_function_space():
+            raise ValueError("Element mismatch between new and old arguments (trial functions).")
+        if new_v.ufl_function_space() != old_v.ufl_function_space():
+            raise ValueError("Element mismatch between new and old arguments (test functions).")
+
+        if check_parts and (new_u.part() != old_v.part() or new_v.part() != old_u.part()):
+            raise ValueError("Ordering of new arguments is the same as the old arguments!")
+
+    if not is_mixed:
+        if len(arguments) != 2:
+            raise ValueError("Expecting bilinear form.")
+
+        v, u = arguments
+        if v.number() >= u.number():
+            raise ValueError("Mistaken assumption in code!")
+        if reordered_arguments is None:
+            assert u.part() is None and v.part() is None
+            new_u = Argument(u.ufl_function_space(), number=v.number())
+            new_v = Argument(v.ufl_function_space(), number=u.number())
+        else:
+            assert isinstance(reordered_arguments, tuple) and len(reordered_arguments) == 2
+            u_arg, v_arg = reordered_arguments[0], reordered_arguments[1]
+            assert isinstance(u_arg, Argument) and isinstance(v_arg, Argument)
+            new_u, new_v = u_arg, v_arg
+
+        validate_mapping(v, u, new_v, new_u, check_parts=True)
+
+        return map_integrands(Conj, replace(form, {v: new_v, u: new_u}))
+    else:
+        form_blocked = extract_blocks(form, arity=2)
+        # Apply mapping block-by-block and sum
+        form_adj = 0
+        assert isinstance(form_blocked, tuple)
+        for i, row in enumerate(form_blocked):
+            assert isinstance(row, tuple)
+            for j, block in enumerate(row):
+                if block is not None:
+                    v, u = block.arguments()
+                    if reordered_arguments is not None:
+                        new_v = reordered_arguments[i][1]
+                        new_u = reordered_arguments[j][0]
+                    else:
+                        new_v = Argument(v.ufl_function_space(), number=u.number(), part=v.part())
+                        new_u = Argument(u.ufl_function_space(), number=v.number(), part=u.part())
+                    local_map = {v: new_v, u: new_u}
+                    validate_mapping(v, u, new_v, new_u)
+                    form_adj += map_integrands(Conj, replace(block, local_map))
+
+        return form_adj
diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py
index 68f5aaf..2f9386b 100644
--- a/src/dolfinx_adjoint/solvers.py
+++ b/src/dolfinx_adjoint/solvers.py
@@ -153,7 +153,7 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ
                 self._lhs,  # type: ignore[arg-type]
                 self._rhs,  # type: ignore[arg-type]
                 bcs=self.bcs,
-                u=self.u,
+                u=self.u,  # type: ignore[arg-type]
                 P=self._preconditioner,  # type: ignore[arg-type]
                 kind=self._kind,  # type: ignore[arg-type]
                 petsc_options=self._petsc_options,
@@ -292,10 +292,10 @@ def solve(self, annotate: bool = True) -> typing.Union[dolfinx.fem.Function, typ
         annotate = pyadjoint.annotate_tape({"annotate": annotate})
         if annotate:
             block = NonlinearProblemBlock(
-                J=self._lhs,
+                J=self._lhs,  # type: ignore[arg-type]
                 F=self._rhs,  # type: ignore[arg-type]
                 bcs=self._bcs,
-                u=self.u,
+                u=self.u,  # type: ignore[arg-type]
                 P=self._preconditioner,  # type: ignore[arg-type]
                 kind=self._kind,  # type: ignore[arg-type]
                 petsc_options=self._petsc_options,
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index f10814d..ef1c6a3 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -14,9 +14,9 @@
 )
 from pyadjoint.tape import no_annotations
 
+from ..blocks._vector import _SpecialVector, _vector
 from ..blocks.assembly import assemble_compiled_form
 from ..utils import function_from_vector, gather
-from ..blocks._vector import _vector, _SpecialVector
 
 
 def _create_function(
@@ -87,7 +87,7 @@ def _ad_init_object(cls, obj):
         return cls(obj.function_space, obj.x, obj.name)
 
     @property
-    def index_map(self) -> dolfinx.cpp.la.IndexMap:
+    def index_map(self) -> dolfinx.cpp.la.IndexMap:  # type: ignore
         """Return the index map of the function's vector."""
         return self.x.index_map
 
diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py
index 5b5730e..6b0b01f 100644
--- a/src/dolfinx_adjoint/utils.py
+++ b/src/dolfinx_adjoint/utils.py
@@ -59,6 +59,7 @@ def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.F
     function.x.array[:] = 0.0
     floatifier = Floatify()
     for weight, func in pairs:
+        assert isinstance(func, dolfinx.fem.Function), "All operands in the linear combination must be Functions."
         if not func.function_space == function.function_space:
             raise ValueError("Function spaces of all functions in the linear combination must match for assignment.")
         function.x.array[:] += floatifier.process(weight) * func.x.array[:]
diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
new file mode 100644
index 0000000..940cc2f
--- /dev/null
+++ b/tests/test_blocked_problem.py
@@ -0,0 +1,122 @@
+from mpi4py import MPI
+
+import basix.ufl
+import dolfinx
+import numpy as np
+import pyadjoint
+import pytest
+import ufl
+
+from dolfinx_adjoint import Function, assemble_scalar
+from dolfinx_adjoint.solvers import LinearProblem
+
+
+@pytest.fixture(scope="module")
+def mesh_2D():
+    return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7)
+
+
+@pytest.mark.parametrize("use_mixed_space", [True, False])
+def test_solver(use_mixed_space: bool, mesh_2D):
+    pyadjoint.get_working_tape().clear_tape()
+    mesh = mesh_2D
+    el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,))
+    el_p = basix.ufl.element("P", mesh.basix_cell(), 1)
+    V = dolfinx.fem.functionspace(mesh, el_u)
+    Q = dolfinx.fem.functionspace(mesh, el_p)
+    dx = ufl.Measure("dx", domain=mesh)
+
+    def a00(u, v):
+        return ufl.inner(ufl.grad(u), ufl.grad(v)) * dx
+
+    def a01(p, v):
+        return ufl.inner(p, ufl.div(v)) * dx
+
+    def a10(q, u):
+        return ufl.inner(q, ufl.div(u)) * dx
+
+    def L0(f, v):
+        return ufl.inner(f, v) * dx
+
+    def L1(mesh, q):
+        return dolfinx.fem.Constant(mesh, 0.0) * q * dx
+
+    Z = dolfinx.fem.functionspace(mesh, ("DG", 0, (mesh.geometry.dim,)))
+    f = Function(Z, name="control")
+    f.interpolate(lambda x: (np.sin(x[0]), x[1]))
+
+    if use_mixed_space:
+        W = ufl.MixedFunctionSpace(*[V, Q])
+        u, p = ufl.TrialFunctions(W)
+        v, q = ufl.TestFunctions(W)
+        a = ufl.extract_blocks(a00(u, v) + a01(p, v) + a10(q, u))
+        L = ufl.extract_blocks(L0(f, v) + L1(mesh, q))
+    else:
+        u = ufl.TrialFunction(V)
+        p = ufl.TrialFunction(Q)
+        v = ufl.TestFunction(V)
+        q = ufl.TestFunction(Q)
+        a = [[a00(u, v), a01(p, v)], [a10(q, u), None]]
+        L = [L0(f, v), L1(mesh, q)]
+
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+    boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)
+    bc_val = dolfinx.fem.Constant(mesh, np.zeros((mesh.geometry.dim,), dtype=dolfinx.default_scalar_type))
+    bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
+
+    options = {
+        "ksp_monitor": None,
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+    }
+    uh, ph = (Function(V, name="state"), Function(Q, name="pressure"))
+    problem = LinearProblem(
+        a,
+        L,
+        u=[uh, ph],
+        bcs=[bc],
+        petsc_options=options,
+        adjoint_petsc_options=options,
+        tlm_petsc_options=options,
+    )
+    problem.solve()
+
+    x = ufl.SpatialCoordinate(mesh)
+    c = ufl.as_vector((2 * ufl.sin(x[0]), 3 * ufl.cos(x[1])))
+    error = ufl.inner(uh - c, uh - c) * ufl.inner(uh - c, uh - c) * ufl.dx
+    J = assemble_scalar(error)
+
+    control = pyadjoint.Control(f)
+    Jh = pyadjoint.ReducedFunctional(J, control)
+    d = Function(Z)
+    d.interpolate(lambda x: (10 * x[0], x[1]))
+
+    e = Function(Z)
+    e.interpolate(lambda x: (1e3 * np.sin(x[1]), 1e3 * x[0]))
+    min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=0)
+    assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 1.0, got {min_rate}"
+
+    Jh.derivative()
+    min_rate = pyadjoint.taylor_test(Jh, d, e)
+    assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2), f"Expected convergence rate close to 2.0, got {min_rate}"
+
+    Jh(d)
+    dJdm = Jh.derivative()._ad_dot(e)
+    hessian = Jh.hessian(e)
+    dHddu = hessian._ad_dot(e)
+    min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu)
+    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+
+    z = Function(Z)
+    z.interpolate(lambda x: (np.sin(x[1]), -(x[0] ** 2)))
+    f = Function(Z)
+    f.interpolate(lambda x: (1e4 * x[0], 1e5 * np.sin(x[1])))
+    Jh(z)
+    dJdm = Jh.derivative()._ad_dot(f)
+    hessian = Jh.hessian(f)
+    dHddu = hessian._ad_dot(f)
+    min_rate = pyadjoint.taylor_test(Jh, z, f, dJdm=dJdm, Hm=dHddu)
+    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py
index 7660efc..0687a11 100644
--- a/tests/test_dirichlet_bc.py
+++ b/tests/test_dirichlet_bc.py
@@ -118,8 +118,14 @@ def test_time_dependent_bc_replay():
 
     # Use native dolfinx here! PyAdjoint traces the bc_func inside it.
     bc = dirichletbc(bc_func, boundary_dofs)
-
-    problem = LinearProblem(a, L, bcs=[bc], u=uh)
+    petsc_options = {
+        "ksp_monitor": None,
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+    }
+    problem = LinearProblem(a, L, bcs=[bc], u=uh, petsc_options=petsc_options)
 
     J = 0.0
 
diff --git a/tests/test_linear_solver.py b/tests/test_linear_solver.py
index 40d54cf..09cb830 100644
--- a/tests/test_linear_solver.py
+++ b/tests/test_linear_solver.py
@@ -88,3 +88,77 @@ def test_solver(mesh_var_name: str, request, constant: typing.Union[float, int,
     dHddu = hessian._ad_dot(e)
     min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=dJdm, Hm=dHddu)
     assert np.isclose(min_rate, 3.0, rtol=5e-3, atol=5e-3), f"Expected convergence rate close to 3.0, got {min_rate}"
+
+
+def test_linear_mixed_derivative_hessian(mesh_2D):
+    """Test LinearProblem with a control in the bilinear form (d2F/dudm != 0)."""
+    pyadjoint.get_working_tape().clear_tape()
+    mesh = mesh_2D
+
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    uh = Function(V, name="state")
+    v = ufl.TestFunction(V)
+    u_trial = ufl.TrialFunction(V)
+
+    # Control variable (conductivity)
+    m = Function(V, name="control")
+    m.interpolate(lambda x: 1.0 + x[0] ** 2 + x[1] ** 2)
+
+    # Constant Source term
+    f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0))
+
+    # Linear problem where the bilinear form explicitly depends on the control 'm'
+    a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx
+    L = f * v * ufl.dx
+
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+    boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)
+    bc = dolfinx.fem.dirichletbc(dolfinx.default_scalar_type(0.0), boundary_dofs, V)
+
+    # Forward Solve
+    petsc_options = {
+        "ksp_monitor": None,
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+    }
+    problem = LinearProblem(
+        a,
+        L,
+        bcs=[bc],
+        u=uh,
+        petsc_options=petsc_options,
+        adjoint_petsc_options=petsc_options,
+        tlm_petsc_options=petsc_options,
+    )
+    problem.solve()
+
+    # Define Objective
+    d = Function(V)
+    d.interpolate(lambda x: np.sin(np.pi * x[0]))
+    J = assemble_scalar(0.5 * ufl.inner(uh - d, uh - d) * ufl.dx)
+
+    control = pyadjoint.Control(m)
+    Jh = pyadjoint.ReducedFunctional(J, control)
+
+    # Perturbation
+    dm = Function(V)
+    dm.interpolate(lambda x: 2 * np.sin(x[0] * np.pi) * np.cos(x[1] * np.pi))
+
+    # Perturbation test
+    min_rate = pyadjoint.taylor_test(Jh, m, dm, dJdm=0)
+    assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 1.0, got {min_rate}"
+    Jh(m)
+
+    # Gradient Taylor Test
+    min_rate_grad = pyadjoint.taylor_test(Jh, m, dm)
+    assert np.isclose(min_rate_grad, 2.0, rtol=1e-1, atol=1e-1), f"Grad rate failed: {min_rate_grad}"
+
+    # Hessian Taylor Test
+    Jh(m)
+    dJ = Jh.derivative()._ad_dot(dm)
+    H = Jh.hessian(dm)._ad_dot(dm)
+    min_rate_hess = pyadjoint.taylor_test(Jh, m, dm, dJdm=dJ, Hm=H)
+    assert np.isclose(min_rate_hess, 3.0, rtol=1e-1, atol=1e-1), f"Hessian rate failed: {min_rate_hess}"
diff --git a/tests/test_nonlinear_problem.py b/tests/test_nonlinear_problem.py
new file mode 100644
index 0000000..7d8e37d
--- /dev/null
+++ b/tests/test_nonlinear_problem.py
@@ -0,0 +1,94 @@
+from mpi4py import MPI
+
+import dolfinx
+import numpy as np
+import pyadjoint
+import ufl
+
+from dolfinx_adjoint import Function, assemble_scalar
+from dolfinx_adjoint.solvers import NonlinearProblem
+
+
+def test_sequential_nonlinear_problems():
+    """
+    Test two cascaded non-linear PDEs.
+    PDE 1: -div(u1 * grad(u1)) = f
+    PDE 2: -div(u2 * grad(u2)) = u1
+    """
+    pyadjoint.get_working_tape().clear_tape()
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7)
+
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+
+    f = Function(V, name="control")
+    # Keep the control positive to avoid degeneracy in the diffusion tensor
+    f.interpolate(lambda x: 2.0 + np.sin(x[0]))
+
+    u1 = Function(V, name="state_1")
+    u1.interpolate(lambda x: np.ones_like(x[0]))  # Non-zero initial guess
+    v1 = ufl.TestFunction(V)
+
+    F1 = (1 + u1**2) * ufl.inner(ufl.grad(u1), ufl.grad(v1)) * ufl.dx(domain=mesh) - f * v1 * ufl.dx(domain=mesh)
+
+    # Setup PDE 2
+    u2 = Function(V, name="state_2")
+    u2.interpolate(lambda x: np.ones_like(x[0]))  # Non-zero initial guess
+    v2 = ufl.TestFunction(V)
+    F2 = (2 + u2**2) * ufl.inner(ufl.grad(u2), ufl.grad(v2)) * ufl.dx(domain=mesh) - u1 * v2 * ufl.dx(domain=mesh)
+
+    # 4. Boundary Conditions (u = 1.0 on boundary)
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+    boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets)
+    bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))
+    bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
+
+    # Use SNES options for the nonlinear solver
+    options = {
+        "snes_monitor": None,
+        "snes_error_if_not_converged": True,
+        "snes_type": "newtonls",
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "pc_factor_mat_solver_type": "mumps",
+    }
+
+    # 5. Solve the Cascade
+    problem1 = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    problem1.solve()
+
+    problem2 = NonlinearProblem(F2, u=u2, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    problem2.solve()
+
+    # 6. Objective (using the cubed error to ensure a 3.0 Hessian rate)
+    d = pyadjoint.AdjFloat(0.2)
+    error = (u2 - d) ** 3 * ufl.dx(domain=mesh)
+    J = assemble_scalar(error)
+
+    # 7. Taylor Tests
+    control = pyadjoint.Control(f)
+    Jh = pyadjoint.ReducedFunctional(J, control)
+
+    ctrl_eval = Function(V)
+    ctrl_eval.interpolate(lambda x: 4.0 + np.sin(x[0]))
+
+    pert = Function(V)
+    pert.interpolate(lambda x: 15.1 * np.cos(x[1]))
+
+    Jh(ctrl_eval)
+    min_rate_grad = pyadjoint.taylor_test(Jh, ctrl_eval, pert, dJdm=0)
+    print("\n--- 1st-order Taylor test ---")
+    Jh(ctrl_eval)
+    min_rate_grad = pyadjoint.taylor_test(Jh, ctrl_eval, pert)
+    assert np.isclose(min_rate_grad, 2.0, rtol=1e-2, atol=1e-2), f"Expected 2.0, got {min_rate_grad}"
+
+    print("\n--- 2nd-order Taylor test ---")
+    Jh(ctrl_eval)
+    dJdm = Jh.derivative()._ad_dot(pert)
+    dHddu = Jh.hessian(pert)._ad_dot(pert)
+    min_rate_hess = pyadjoint.taylor_test(Jh, ctrl_eval, pert, dJdm=dJdm, Hm=dHddu)
+    assert np.isclose(min_rate_hess, 3.0, rtol=1e-2, atol=1e-2), f"Expected 3.0, got {min_rate_hess}"
+
+
+if __name__ == "__main__":
+    test_sequential_nonlinear_problems()
diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
new file mode 100644
index 0000000..9577309
--- /dev/null
+++ b/tests/test_tlm_update.py
@@ -0,0 +1,175 @@
+"""Regression tests for the tangent-linear solver cached on ``LinearProblemBlock``.
+
+``construct_tlm_solver`` compiles ``_compute_residual_derivative()`` once and caches the
+resulting solver on the block.  That form is built from ``block_variable.saved_output``,
+which is a fresh object after every control update, so the cached operator can stay
+pinned to the evaluation point it was first built at.  When the control appears in the
+bilinear form -- so that dF/du genuinely depends on it -- every Hessian computed after
+the first one is then silently wrong.
+"""
+
+from mpi4py import MPI
+
+import basix.ufl
+import dolfinx
+import numpy as np
+import pyadjoint
+import pytest
+import ufl
+
+from dolfinx_adjoint import Function, assemble_scalar
+from dolfinx_adjoint.solvers import LinearProblem
+
+direct_solve = {
+    "ksp_type": "preonly",
+    "pc_type": "lu",
+    "ksp_error_if_not_converged": True,
+    "pc_factor_mat_solver_type": "mumps",
+}
+
+
+@pytest.fixture(scope="module")
+def mesh_2D():
+    return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7)
+
+
+def _viscous_stokes(mesh):
+    """Blocked Stokes-like problem whose control ``mu`` sits inside ``a[0][0]``.
+
+    The control has to enter the bilinear form for this test to have any teeth: if it
+    only enters ``L``, dF/du is independent of the control and a stale tangent-linear
+    operator is indistinguishable from a fresh one.
+    """
+    el_u = basix.ufl.element("P", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,))
+    el_p = basix.ufl.element("P", mesh.basix_cell(), 1)
+    V = dolfinx.fem.functionspace(mesh, el_u)
+    Q = dolfinx.fem.functionspace(mesh, el_p)
+    Z = dolfinx.fem.functionspace(mesh, ("DG", 0))
+    dx = ufl.Measure("dx", domain=mesh)
+
+    mu = Function(Z, name="viscosity")
+    mu.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0]))
+
+    u, p = ufl.TrialFunction(V), ufl.TrialFunction(Q)
+    v, q = ufl.TestFunction(V), ufl.TestFunction(Q)
+    # A rotational (non-conservative) body force.  A constant force would be balanced
+    # exactly by the pressure in a closed incompressible box, leaving u == 0 and making
+    # the functional independent of the control.  The 1e3 only sets the scale of the
+    # state, which keeps the Taylor remainders clear of pyadjoint's absolute
+    # machine-precision warning threshold.
+    x = ufl.SpatialCoordinate(mesh)
+    f = 1e3 * ufl.as_vector((ufl.sin(ufl.pi * x[1]), ufl.cos(ufl.pi * x[0])))
+
+    a = [
+        [ufl.inner(mu * ufl.grad(u), ufl.grad(v)) * dx, ufl.inner(p, ufl.div(v)) * dx],
+        [ufl.inner(q, ufl.div(u)) * dx, None],
+    ]
+    L = [ufl.inner(f, v) * dx, dolfinx.fem.Constant(mesh, 0.0) * q * dx]
+
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+    dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets)
+    zero = dolfinx.fem.Constant(mesh, np.zeros(mesh.geometry.dim, dtype=dolfinx.default_scalar_type))
+    bc = dolfinx.fem.dirichletbc(zero, dofs, V)
+
+    uh, ph = Function(V, name="velocity"), Function(Q, name="pressure")
+    problem = LinearProblem(
+        a,
+        L,
+        u=[uh, ph],
+        bcs=[bc],
+        petsc_options=direct_solve,
+        adjoint_petsc_options=direct_solve,
+        tlm_petsc_options=direct_solve,
+    )
+    problem.solve()
+
+    # Quartic in the state and with no constant offset, so that the Taylor remainders
+    # stay well above round-off: a functional dominated by a control-independent term
+    # bottoms out at machine precision before the third-order rate is visible.
+    J = assemble_scalar(ufl.inner(uh, uh) ** 2 * dx)
+    return pyadjoint.ReducedFunctional(J, pyadjoint.Control(mu)), Z
+
+
+@pytest.mark.parametrize("warm_up_at_another_point", [False, True])
+def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another_point, mesh_2D):
+    """The Hessian at ``m2`` must not depend on whether ``J`` was evaluated at ``m1`` first.
+
+    Both parametrizations run the identical second-order Taylor test at ``m2``.  The only
+    difference is a prior evaluate-and-differentiate sweep at a *different* control value,
+    which must not change the answer.
+    """
+    pyadjoint.get_working_tape().clear_tape()
+    Jh, Z = _viscous_stokes(mesh_2D)
+
+    m1, m2 = Function(Z), Function(Z)
+    m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1]))
+    h = Function(Z)
+    h.interpolate(lambda x: 10.0 + 3.2 * np.sin(3 * x[0]))
+
+    if warm_up_at_another_point:
+        Jh(m1)
+        Jh.derivative()
+        Jh.hessian(h)
+
+    Jh(m2)
+    dJdm = Jh.derivative()._ad_dot(h)
+    Hm = Jh.hessian(h)._ad_dot(h)
+
+    min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm)
+    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected convergence rate close to 3.0, got {min_rate}"
+
+
+def test_hessian_mpi_breakdown(mesh_2D):
+    pyadjoint.get_working_tape().clear_tape()
+    Jh, Z = _viscous_stokes(mesh_2D)
+
+    m1, m2 = Function(Z), Function(Z)
+    m1.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1]))
+    h = Function(Z)
+    h.interpolate(lambda x: 3.0 + 2.0 * np.sin(3 * x[0]))
+
+    # === COLD START ===
+    J_cold = float(Jh(m2))
+    dJ_cold = Jh.derivative()
+    H_cold = Jh.hessian(h)
+
+    # Extract underlying FEniCSx arrays securely
+    dJ_cold_array = dJ_cold.x.array.copy() if hasattr(dJ_cold, "x") else dJ_cold.array.copy()
+    H_cold_array = H_cold.x.array.copy() if hasattr(H_cold, "x") else H_cold.array.copy()
+
+    # === WARM START (Pollution check) ===
+    Jh(m1)
+    Jh.derivative()
+    Jh.hessian(h)
+
+    J_warm = float(Jh(m2))
+    dJ_warm = Jh.derivative()
+    H_warm = Jh.hessian(h)
+
+    dJ_warm_array = dJ_warm.x.array if hasattr(dJ_warm, "x") else dJ_warm.array
+    H_warm_array = H_warm.x.array if hasattr(H_warm, "x") else H_warm.array
+
+    # === ASSERTIONS ===
+
+    comm = mesh_2D.comm
+    rank = comm.rank
+
+    # 1. Forward State Check
+    assert np.isclose(J_cold, J_warm), f"Rank {rank}: Forward evaluation J(m2) differs!"
+
+    # 2. Gradient Check
+    grad_diff = np.linalg.norm(dJ_cold_array - dJ_warm_array)
+    assert grad_diff < 1e-10, f"Rank {rank}: Gradient differs after warm up! Diff: {grad_diff}"
+
+    # 3. Hessian Check
+    hess_diff = np.linalg.norm(H_cold_array - H_warm_array)
+    assert hess_diff < 1e-10, f"Rank {rank}: Hessian differs after warm up! Diff: {hess_diff}"
+
+    # 4. If arrays match, run Taylor test
+    dJdm = dJ_warm._ad_dot(h)
+    Hm = H_warm._ad_dot(h)
+    min_rate = pyadjoint.taylor_test(Jh, m2, h, dJdm=dJdm, Hm=Hm)
+    assert np.isclose(min_rate, 3.0, rtol=0.1, atol=0.1), f"Expected 3.0, got {min_rate}"