From bbb4567bb0756124ff2ac74711292657980c4641 Mon Sep 17 00:00:00 2001
From: jorgensd <dokken92@gmail.com>
Date: Tue, 1 Sep 2026 21:46:33 +0000
Subject: [PATCH 01/22] Mypy fixes

---
 src/dolfinx_adjoint/blocks/solvers.py | 22 ++++++++++++-
 src/dolfinx_adjoint/solvers.py        | 26 +++++++--------
 src/dolfinx_adjoint/ufl_utils.py      | 46 ++++++++++-----------------
 3 files changed, 50 insertions(+), 44 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py
index 62924b3..eeeb5a2 100644
--- a/src/dolfinx_adjoint/blocks/solvers.py
+++ b/src/dolfinx_adjoint/blocks/solvers.py
@@ -15,13 +15,33 @@
 
 from ..types import Function
 from ..typing_utils import NestedMutableSequence
-from ..ufl_utils import assign_mixed_parts, collect_coefficients, sum_form
+from ..ufl_utils import assign_mixed_parts, sum_form
 from .assembly import _create_vector, _SpecialVector, assemble_compiled_form
 
 if typing.TYPE_CHECKING:
     from ..solvers import LinearProblem, NonlinearProblem
 
 
+def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set[Function]:
+    """Return the set of UFL coefficients appearing anywhere in ``form``.
+
+    ``form`` may be a single form or an arbitrarily nested sequence of forms
+    (entries may be ``None``, e.g. a zero block in a blocked system). Plain set
+    union rather than ``sum_form``: unlike summing, this never requires the
+    sub-forms' arguments to be mutually compatible (e.g. carry matching
+    ``part()`` tags), which a blocked ``NonlinearProblem``'s forms are not
+    required to be before ``assign_mixed_parts`` runs.
+    """
+    if form is None:
+        return set()
+    if isinstance(form, ufl.Form):
+        return set(form.coefficients())
+    coefficients: set = set()
+    for f in form:
+        coefficients |= collect_coefficients(f)
+    return coefficients
+
+
 def _map_block_variables_to_form(
     form: ufl.Form | NestedMutableSequence[ufl.Form] | None,
     block_variables: typing.Iterable[pyadjoint.block_variable.BlockVariable],
diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py
index b49a639..575ec59 100644
--- a/src/dolfinx_adjoint/solvers.py
+++ b/src/dolfinx_adjoint/solvers.py
@@ -8,13 +8,12 @@
 import pyadjoint
 import ufl
 from dolfinx.fem.function import Function as _Function
-
-from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase
+from .typing_utils import NestedSequence
+from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase, collect_coefficients
 from .petsc_utils import HomogeneousBCLinearProblem
 from .types import Function
 from .ufl_utils import (
     assign_mixed_parts,
-    collect_coefficients,
     compute_adjoint,
     get_sorted_arguments,
     recursive_replace,
@@ -107,13 +106,13 @@ class HessianTemplates(typing.NamedTuple):
             per-dependency (not per-row) shape as ``fixed``.
     """
 
-    soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form]
+    soa_self: NestedSequence[dolfinx.fem.Form]
     soa_cross: dict
     fixed: dict
     cross: dict
 
 
-def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl.Argument]) -> list[ufl.form.BaseForm]:
+def _pad_blocks_by_part(form: ufl.Form, test_funcs: typing.Sequence[ufl.Argument]) -> list[ufl.Form | ufl.ZeroBaseForm]:
     """Split a blocked one-form into one entry per ``test_funcs`` part, in part order.
 
     {py:func}`ufl.extract_blocks` only returns an entry for a part that actually appears in
@@ -128,9 +127,10 @@ def _pad_blocks_by_part(form: ufl.form.BaseForm, test_funcs: typing.Sequence[ufl
     no parts for {py:func}`ufl.extract_blocks` to find, so every row is padded to zero directly
     instead.
     """
-    padded: list[ufl.form.BaseForm] = [ufl.ZeroBaseForm((test,)) for test in test_funcs]
+    padded: list[ufl.Form | ufl.ZeroBaseForm] = [ufl.ZeroBaseForm((test,)) for test in test_funcs]
     if form.empty():
         return padded
+    assert isinstance(form, ufl.Form)
     for block in ufl.extract_blocks(form):
         args = block.arguments()
         assert len(args) == 1, "Expected a single test function in the block."
@@ -147,7 +147,7 @@ def _build_soa_self_template(
     jit_options: dict | None,
     form_compiler_options: dict | None,
     entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None,
-) -> dolfinx.fem.Form:
+) -> NestedSequence[dolfinx.fem.Form]:
     """Build the SOA self-term ``adjoint(d2F/du2) . adjoint_solution``.
 
     The same computation for both {py:class}`~dolfinx_adjoint.LinearProblem` and
@@ -164,7 +164,7 @@ def _build_soa_self_template(
         soa_self_form = ufl.ZeroBaseForm((dFdu_template.arguments()[0],))
     else:
         soa_self_form = ufl.action(ufl.adjoint(d2Fdu2), adjoint_solution_placeholder)
-    return dolfinx.fem.form(
+    return dolfinx.fem.form(  # type: ignore[call-overload]
         soa_self_form,
         jit_options=jit_options,
         form_compiler_options=form_compiler_options,
@@ -264,7 +264,7 @@ def _init_adjoint_state(self) -> None:
         self._tlm_solver: HomogeneousBCLinearProblem | None = None
         self._residual_template: ufl.Form | None = None
         self._residual_state_placeholder: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None
-        self._dFdu_template: ufl.Form | typing.Sequence | None = None
+        self._dFdu_template: ufl.Form | None = None
         self._dFdu_adj_template: ufl.Form | typing.Sequence | None = None
         self._tlm_rhs_templates: dict | None = None
         self._tlm_seed_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {}
@@ -296,7 +296,7 @@ def _get_or_build_residual_template(
             a blocked problem).
         """
 
-    def _get_or_build_dFdu_template(self) -> ufl.Form | typing.Sequence:
+    def _get_or_build_dFdu_template(self) -> ufl.Form:
         """Build (once) and return dF/du, evaluated at the residual template's state placeholder."""
         # Shared by both classes: derived from _get_or_build_residual_template by
         # symbolic differentiation (free at compile time) rather than
@@ -417,7 +417,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
             assert isinstance(dFdu_adj_template, ufl.Form)
 
             blocked = isinstance(self._u, list)
-            soa_self: dolfinx.fem.Form | list[dolfinx.fem.Form]
+            soa_self: NestedSequence[dolfinx.fem.Form]
             if blocked:
                 # One placeholder Function per output block, mirroring how
                 # _get_or_build_hessian_templates's scalar branch below uses a
@@ -451,7 +451,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
                 # happened to eliminate entirely.
                 soa_self = [
                     dolfinx.fem.form(
-                        form_i,
+                        form_i,  # type: ignore[arg-type]
                         jit_options=self._jit_options,
                         form_compiler_options=self._form_compiler_options,
                         entity_maps=self._entity_maps,
@@ -507,7 +507,7 @@ def _get_or_build_hessian_templates(self) -> HessianTemplates:
                     if blocked:
                         soa_cross_templates[c] = [
                             dolfinx.fem.form(
-                                form_i,
+                                form_i,  # type: ignore[arg-type]
                                 jit_options=self._jit_options,
                                 form_compiler_options=self._form_compiler_options,
                                 entity_maps=self._entity_maps,
diff --git a/src/dolfinx_adjoint/ufl_utils.py b/src/dolfinx_adjoint/ufl_utils.py
index bcd6551..60e8005 100644
--- a/src/dolfinx_adjoint/ufl_utils.py
+++ b/src/dolfinx_adjoint/ufl_utils.py
@@ -9,7 +9,7 @@
 
 
 def recursive_space_discovery(
-    obj: NestedSequence[ufl.Form], indices: tuple[int, ...], spaces: dict[int, ufl.FunctionSpace]
+    obj: NestedSequence[ufl.Form | None], indices: tuple[int, ...], spaces: dict[int, ufl.FunctionSpace]
 ) -> None:
     """Recursively discover, for each row/column index, the function space of the
     (as yet unassigned) argument occupying that position.
@@ -34,15 +34,17 @@ def recursive_space_discovery(
         for i, item in enumerate(obj):
             if item is not None:
                 recursive_space_discovery(item, indices + (i,), spaces)
+    elif obj is None:
+        return
     else:
         raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}")
 
 
 def build_argument_replacement_map(
-    obj: NestedSequence[ufl.Form],
+    obj: NestedSequence[ufl.Form | None],
     indices: tuple[int, ...],
-    test_functions: typing.Sequence[ufl.TestFunction],
-    trial_functions: typing.Sequence[ufl.TrialFunction],
+    test_functions: typing.Sequence[ufl.Argument],
+    trial_functions: typing.Sequence[ufl.Argument],
     replace_map: dict[ufl.Argument, ufl.Argument],
 ) -> None:
     """
@@ -66,17 +68,21 @@ def build_argument_replacement_map(
         for i, item in enumerate(obj):
             if item is not None:
                 build_argument_replacement_map(item, indices + (i,), test_functions, trial_functions, replace_map)
+    elif obj is None:
+        return
+    else:
+        raise TypeError(f"Expected ufl.Form or iterable, got {type(obj)}")
 
 
 @typing.overload
-def assign_mixed_parts[T: NestedSequence[ufl.Form]](form1: T, /) -> T: ...
+def assign_mixed_parts[T: NestedSequence[ufl.Form | None]](form1: T, /) -> T: ...
 @typing.overload
-def assign_mixed_parts[T: NestedSequence[ufl.Form], S: NestedSequence[ufl.Form]](
+def assign_mixed_parts[T: NestedSequence[ufl.Form | None], S: NestedSequence[ufl.Form | None]](
     form1: T, form2: S, /
 ) -> tuple[T, S]: ...
 def assign_mixed_parts(
-    *form_structs: NestedSequence[ufl.Form],
-) -> NestedSequence[ufl.Form] | tuple[NestedSequence[ufl.Form], ...]:
+    *form_structs: NestedSequence[ufl.Form | None],
+) -> NestedSequence[ufl.Form | None] | tuple[NestedSequence[ufl.Form | None], ...]:
     """
     Recursively assigns mixed-space `part` indices to {py:class}`ufl.Argument`
     (test and trial functions), within nested iterables of forms.
@@ -106,7 +112,7 @@ def assign_mixed_parts(
         {py:class}`ufl.MixedFunctionSpace` built from the row/column function spaces
         discovered while walking the structure.
     """
-    spaces: dict[int, ufl.functionspace.AbstractFunctionSpace] = {}
+    spaces: dict[int, ufl.FunctionSpace] = {}
     for struct in form_structs:
         recursive_space_discovery(struct, (), spaces)
 
@@ -133,26 +139,6 @@ def get_sorted_arguments(arguments: typing.Iterable[ufl.Argument], number: int)
     return sorted(filter(lambda x: x.number() == number, arguments), key=lambda a: a.part())
 
 
-def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set[ufl.Coefficient]:
-    """Return the set of UFL coefficients appearing anywhere in ``form``.
-
-    ``form`` may be a single form or an arbitrarily nested sequence of forms
-    (entries may be ``None``, e.g. a zero block in a blocked system). Plain set
-    union rather than ``sum_form``: unlike summing, this never requires the
-    sub-forms' arguments to be mutually compatible (e.g. carry matching
-    ``part()`` tags), which a blocked ``NonlinearProblem``'s forms are not
-    required to be before ``assign_mixed_parts`` runs.
-    """
-    if form is None:
-        return set()
-    if isinstance(form, ufl.Form):
-        return set(form.coefficients())
-    coefficients: set = set()
-    for f in form:
-        coefficients |= collect_coefficients(f)
-    return coefficients
-
-
 def sum_form(form: NestedSequence[ufl.Form | None]) -> ufl.Form | None:
     """Sum a blocked form into a single form."""
     # Handle top-level None
@@ -195,7 +181,7 @@ def compute_adjoint(form: ufl.Form) -> typing.Sequence[typing.Sequence[ufl.Form]
     return ufl.extract_blocks(compute_form_adjoint(form))
 
 
-def recursive_replace(form: ufl.Form | typing.Sequence | None, placeholders: dict) -> ufl.Form | typing.Sequence | None:
+def recursive_replace(form: NestedSequence[ufl.Form | None], placeholders: dict) -> NestedSequence[ufl.Form | None]:
     """Recursively apply {py:func}`ufl.replace` to a (possibly nested) form structure.
 
     Args:

From ac3c7ec10fb2ab27994ca13650753a0b005159e1 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 06:27:14 +0000
Subject: [PATCH 02/22] Ruff formatting

---
 src/dolfinx_adjoint/blocks/function_assigner.py | 1 -
 src/dolfinx_adjoint/solvers.py                  | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py
index 237efc7..e8190eb 100644
--- a/src/dolfinx_adjoint/blocks/function_assigner.py
+++ b/src/dolfinx_adjoint/blocks/function_assigner.py
@@ -1,4 +1,3 @@
-
 import dolfinx
 import numpy as np
 import numpy.typing as npt
diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py
index 575ec59..75d28f6 100644
--- a/src/dolfinx_adjoint/solvers.py
+++ b/src/dolfinx_adjoint/solvers.py
@@ -8,10 +8,11 @@
 import pyadjoint
 import ufl
 from dolfinx.fem.function import Function as _Function
-from .typing_utils import NestedSequence
+
 from .blocks.solvers import LinearProblemBlock, NonlinearProblemBlock, _ProblemBlockBase, collect_coefficients
 from .petsc_utils import HomogeneousBCLinearProblem
 from .types import Function
+from .typing_utils import NestedSequence
 from .ufl_utils import (
     assign_mixed_parts,
     compute_adjoint,

From 7bf2a716a9e9523b746f9c3fdd1457de4d742d9c Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 07:00:08 +0000
Subject: [PATCH 03/22] Add stable to test matrix

---
 .github/workflows/test_package.yml | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test_package.yml b/.github/workflows/test_package.yml
index c2ecc24..800baa0 100644
--- a/.github/workflows/test_package.yml
+++ b/.github/workflows/test_package.yml
@@ -15,7 +15,15 @@ on:
 jobs:
   test-code:
     runs-on: ubuntu-latest
-    container: ghcr.io/fenics/dolfinx/dolfinx:nightly
+    container: ghcr.io/fenics/dolfinx/dolfinx:${{ matrix.label }}
+    strategy:
+      fail-fast: false
+      matrix:
+        label: [
+          "stable",
+          "nightly"
+        ]
+
     env:
       OMPI_ALLOW_RUN_AS_ROOT: 1
       OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1

From e3a4c523d0044fa399c4c0ee81812d104a5665b0 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 07:17:08 +0000
Subject: [PATCH 04/22] Tigthen type hinting. Split snes and ksp options in all
 tests.

---
 src/dolfinx_adjoint/blocks/solvers.py | 130 ++++++++++++++----------
 tests/test_blocked_problem.py         |   4 +-
 tests/test_nonlinear_problem.py       |  14 ++-
 tests/test_solver_reuse.py            | 139 +++++++++++++++++++++-----
 tests/test_tlm_update.py              |   4 +-
 5 files changed, 201 insertions(+), 90 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py
index eeeb5a2..ddb71bd 100644
--- a/src/dolfinx_adjoint/blocks/solvers.py
+++ b/src/dolfinx_adjoint/blocks/solvers.py
@@ -11,10 +11,9 @@
 import numpy as np
 import pyadjoint
 import ufl
-from dolfinx.fem.function import Function as _Function
 
 from ..types import Function
-from ..typing_utils import NestedMutableSequence
+from ..typing_utils import NestedSequence
 from ..ufl_utils import assign_mixed_parts, sum_form
 from .assembly import _create_vector, _SpecialVector, assemble_compiled_form
 
@@ -43,7 +42,7 @@ def collect_coefficients(form: ufl.Form | typing.Sequence | None) -> set[Functio
 
 
 def _map_block_variables_to_form(
-    form: ufl.Form | NestedMutableSequence[ufl.Form] | None,
+    form: NestedSequence[ufl.Form | None],
     block_variables: typing.Iterable[pyadjoint.block_variable.BlockVariable],
 ) -> dict[Function, Function]:
     """Map each ``block_variable``'s output coefficient, where it appears in ``form``, to its
@@ -95,10 +94,10 @@ class _ProblemBlockBase(pyadjoint.Block, abc.ABC):
     _problem_ref: weakref.ReferenceType["LinearProblem | NonlinearProblem"]
     _rebuilt_problem: "LinearProblem | NonlinearProblem | None" = None
     _bcs: typing.Sequence[dolfinx.fem.DirichletBC]
-    _u: _Function | typing.Sequence[_Function]
-    _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]
+    _u: Function | typing.Sequence[Function]
+    _adjoint_solutions: Function | typing.Sequence[Function]
+    _second_adjoint_solutions: Function | typing.Sequence[Function]
+    _tlm_solutions: Function | typing.Sequence[Function]
     _jit_options: dict | None
     _form_compiler_options: dict | None
     _entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None
@@ -179,7 +178,7 @@ def _refresh_dFdu_state(self, problem: "LinearProblem | NonlinearProblem") -> No
         """
         pass
 
-    def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] | None) -> dict[Function, Function]:
+    def _create_replace_map(self, form: NestedSequence[ufl.Form | None]) -> dict[Function, Function]:
         """Map each dependency and output to its checkpointed value, wherever it appears in ``form``.
 
         Args:
@@ -195,9 +194,7 @@ def _create_replace_map(self, form: ufl.Form | NestedMutableSequence[ufl.Form] |
         replace_map.update(_map_block_variables_to_form(form, self.get_outputs()))
         return replace_map
 
-    def prepare_evaluate_tlm(
-        self, inputs, tlm_inputs, relevant_outputs
-    ) -> typing.Sequence[Function] | dolfinx.fem.Function:
+    def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs) -> NestedSequence[Function]:
         """Assemble and solve the tangent-linear (TLM) system for this block.
 
         The TLM solver -- and the compiled LHS it solves with, shared verbatim with
@@ -282,7 +279,7 @@ def prepare_evaluate_tlm(
 
         return self._tlm_solutions
 
-    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None) -> dolfinx.fem.Function:
+    def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx: int, prepared=None) -> Function:
         """Return this output's share of the tangent-linear solution already computed
         by {py:meth}`~dolfinx_adjoint.blocks.solvers._ProblemBlockBase.prepare_evaluate_tlm`.
 
@@ -302,7 +299,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         if isinstance(self._tlm_solutions, list):
             return self._tlm_solutions[idx]
         else:
-            assert isinstance(self._tlm_solutions, dolfinx.fem.Function)
+            assert isinstance(self._tlm_solutions, Function)
             return self._tlm_solutions
 
     def prepare_evaluate_adj(
@@ -374,11 +371,6 @@ def prepare_evaluate_adj(
             dolfinx.la.petsc.assign(arrs, dJdu)
             dJdu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)  # type: ignore[arg-type]
 
-        # F_form/replacement_map are still needed by evaluate_adj_component
-        # (to build each dependency's own sensitivity form), but the adjoint
-        # LHS itself is already correct on adjoint_solver -- no rebuild, no
-        # recompile.
-        F_form, replacement_map = self._compute_residual()
         adjoint_solver.solve()
         if isinstance(self._adjoint_solutions, list):
             for adj_sol, sol in zip(self._adjoint_solutions, adjoint_solver.u):
@@ -386,6 +378,12 @@ def prepare_evaluate_adj(
         else:
             assert isinstance(self._adjoint_solutions, dolfinx.fem.Function)
             self._adjoint_solutions.x.array[:] = adjoint_solver.u.x.array[:]
+
+        # F_form/replacement_map are still needed by evaluate_adj_component
+        # (to build each dependency's own sensitivity form), but the adjoint
+        # LHS itself is already correct on adjoint_solver -- no rebuild, no
+        # recompile.
+        F_form, replacement_map = self._compute_residual()
         return F_form, replacement_map
 
     def evaluate_adj_component(
@@ -450,8 +448,12 @@ def evaluate_adj_component(
 
     def prepare_recompute_component(
         self, inputs: typing.Sequence[typing.Any], relevant_outputs: typing.Sequence[typing.Any]
-    ) -> _Function | typing.Sequence[_Function]:
-        """Prepare for recomputing the block with different control inputs, and solve.
+    ) -> Function | typing.Sequence[Function]:
+        """Recompute the block's own forward solution(s) from its checkpointed dependencies and outputs.
+
+        Each problem has replaced its own forms' coefficients with placeholders, which are populated
+        from the block's saved outputs and dependencies here, then the shared forward solver is called
+        once to recompute the solution(s).
 
         The forward solver (``self.get_reference_problem()``) is bound, forever, to
         compiled forms referencing dedicated placeholder coefficients rather
@@ -520,7 +522,7 @@ def recompute_component(
         inputs: typing.Iterable[Function],
         block_variable: pyadjoint.block_variable.BlockVariable,
         idx: int,
-        prepared: _Function | typing.Sequence[_Function],
+        prepared: Function | typing.Sequence[Function],
     ) -> Function:
         """Return an isolated copy of this block's own share of the already-recomputed state.
 
@@ -537,15 +539,27 @@ def recompute_component(
         Returns:
             An isolated copy of this output, so this tape block's own checkpoint
             stays stable even if the shared Problem's unknown is later overwritten
-            by another block's recompute.
+            by another block's recompute. Reuses ``block_variable.checkpoint`` in
+            place when one already exists (mirroring
+            {py:class}`~dolfinx_adjoint.blocks.interpolation.InterpolationBlock`'s
+            recompute and Firedrake's equivalent ``GenericSolveBlock.recompute_component``),
+            since {py:meth}`~dolfinx_adjoint.types.function.Function._ad_create_checkpoint`
+            -- not a bare ``.copy()``, which always returns a plain, non-overloaded
+            ``dolfinx.fem.Function`` regardless of the source's concrete type -- is what
+            correctly builds a *new* one when none exists yet.
         """
-        if isinstance(prepared, dolfinx.fem.Function):
+        if isinstance(prepared, Function):
             assert idx == 0
-            # Return an explicit copy so each tape block gets an isolated state snapshot
-            return prepared.copy()
+            source = prepared
         else:
             assert isinstance(prepared, typing.Sequence)
-            return prepared[idx].copy()
+            source = prepared[idx]
+        checkpoint = block_variable.checkpoint
+        if isinstance(checkpoint, Function):
+            checkpoint.x.array[:] = source.x.array[:]
+            checkpoint.x.scatter_forward()
+            return checkpoint
+        return source._ad_create_checkpoint()
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
         """Assemble and solve the second-order-adjoint (SOA) equation.
@@ -830,9 +844,9 @@ class LinearProblemBlock(_ProblemBlockBase):
     This class extends the `dolfinx.fem.petsc.LinearProblem` to support adjoint methods.
     """
 
-    _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]
+    _adjoint_solutions: Function | typing.Sequence[Function]
+    _tlm_solutions: Function | typing.Sequence[Function]
+    _second_adjoint_solutions: Function | typing.Sequence[Function]
 
     # 2. Overload for the SCALAR case
     @typing.overload
@@ -842,7 +856,7 @@ def __init__(
         L: ufl.Form,
         *,
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: _Function | None = None,
+        u: Function | None = None,
         P: ufl.Form | None = None,
         form_compiler_options: dict | None = None,
         jit_options: dict | None = None,
@@ -863,7 +877,7 @@ def __init__(
         L: typing.Sequence[ufl.Form],
         *,
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: typing.Sequence[_Function] | None = None,
+        u: typing.Sequence[Function] | None = None,
         P: typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         form_compiler_options: dict | None = None,
         jit_options: dict | None = None,
@@ -883,7 +897,7 @@ def __init__(
         L: ufl.Form | typing.Sequence[ufl.Form],
         *,
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: _Function | typing.Sequence[_Function] | None = None,
+        u: Function | typing.Sequence[Function] | None = None,
         P: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         form_compiler_options: dict | None = None,
         jit_options: dict | None = None,
@@ -926,7 +940,7 @@ def __init__(
         self._preconditioner = P
 
         # Create overloaded functions
-        self._u: _Function | typing.Sequence[_Function]
+        self._u: Function | typing.Sequence[Function]
         if isinstance(u, dolfinx.fem.Function):
             self._u = pyadjoint.create_overloaded_object(u)
         elif u is None:
@@ -990,15 +1004,21 @@ def __init__(
         # built once and reused across every block that Problem records
         # instead of once per solve() call.
 
+        # Private, isolated scratch storage for this block's own adjoint/TLM
+        # solutions -- never shared with problem.u or with any other block.
+        # Built via _ad_create_checkpoint(), not a bare .copy(): the latter
+        # always returns a plain, non-overloaded dolfinx.fem.Function
+        # regardless of the source's concrete type (see the same note on
+        # Function._ad_create_checkpoint in types/function.py).
         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()
+            self._adjoint_solutions = self._u._ad_create_checkpoint()
+            self._second_adjoint_solutions = self._u._ad_create_checkpoint()
+            self._tlm_solutions = self._u._ad_create_checkpoint()
         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_solutions = [u._ad_create_checkpoint() for u in self._u]
+            self._second_adjoint_solutions = [u._ad_create_checkpoint() for u in self._u]
+            self._tlm_solutions = [u._ad_create_checkpoint() for u in self._u]
 
     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
@@ -1060,9 +1080,9 @@ class NonlinearProblemBlock(_ProblemBlockBase):
     This class extends the `dolfinx.fem.petsc.NonlinearProblem` to support adjoint methods.
     """
 
-    _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]
+    _adjoint_solutions: Function | typing.Sequence[Function]
+    _second_adjoint_solutions: Function | typing.Sequence[Function]
+    _tlm_solutions: Function | typing.Sequence[Function]
     _rhs: ufl.Form | typing.Sequence[ufl.Form]
 
     @typing.overload
@@ -1070,7 +1090,7 @@ def __init__(
         self,
         F: ufl.Form,
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: dolfinx.fem.Function | None = None,
+        u: Function | None = None,
         J: ufl.Form | None = None,
         P: ufl.Form | None = None,
         form_compiler_options: dict | None = None,
@@ -1090,7 +1110,7 @@ def __init__(
         self,
         F: typing.Sequence[ufl.Form],
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: typing.Sequence[dolfinx.fem.Function] | None = None,
+        u: typing.Sequence[Function] | None = None,
         J: typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         P: typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         form_compiler_options: dict | None = None,
@@ -1109,7 +1129,7 @@ def __init__(
         self,
         F: ufl.Form | typing.Sequence[ufl.Form],
         bcs: typing.Sequence[dolfinx.fem.DirichletBC] | None = None,
-        u: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function] | None = None,
+        u: Function | typing.Sequence[Function] | None = None,
         J: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         P: ufl.Form | typing.Sequence[typing.Sequence[ufl.Form]] | None = None,
         form_compiler_options: dict | None = None,
@@ -1142,7 +1162,7 @@ def __init__(
 
         # Create overloaded functions
         assert u is not None, "Control variable(s) must be provided."
-        self._u: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function]
+        self._u: Function | typing.Sequence[Function]
         if isinstance(u, dolfinx.fem.Function):
             self._u = pyadjoint.create_overloaded_object(u)
             replace_dict = {u: self._u}
@@ -1177,15 +1197,21 @@ def __init__(
         # built once and reused across every block that Problem records
         # instead of once per solve() call.
 
+        # Private, isolated scratch storage for this block's own adjoint/TLM
+        # solutions -- never shared with problem.u or with any other block.
+        # Built via _ad_create_checkpoint(), not a bare .copy(): the latter
+        # always returns a plain, non-overloaded dolfinx.fem.Function
+        # regardless of the source's concrete type (see the same note on
+        # Function._ad_create_checkpoint in types/function.py).
         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]
+            self._adjoint_solutions = self._u._ad_create_checkpoint()
+            self._second_adjoint_solutions = self._u._ad_create_checkpoint()
+            self._tlm_solutions = self._u._ad_create_checkpoint()
         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_solutions = [u._ad_create_checkpoint() for u in self._u]
+            self._second_adjoint_solutions = [u._ad_create_checkpoint() for u in self._u]
+            self._tlm_solutions = [u._ad_create_checkpoint() for u in self._u]
 
     def _compute_residual(self) -> tuple[ufl.Form, dict[Function, Function]]:
         """Build the residual :math:`F(u_b, v) = 0` at the current checkpointed dependency values.
diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index b88740b..97008ad 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -188,10 +188,8 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
         "snes_atol": 1e-9,
         "snes_rtol": 1e-9,
         "snes_stol": 1e-12,
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "pc_factor_mat_solver_type": "mumps",
     }
+    forward_options.update(direct_solve)
     problem = NonlinearProblem(
         F,
         u=[uh, ph],
diff --git a/tests/test_nonlinear_problem.py b/tests/test_nonlinear_problem.py
index 7d8e37d..9cf1528 100644
--- a/tests/test_nonlinear_problem.py
+++ b/tests/test_nonlinear_problem.py
@@ -44,20 +44,24 @@ def test_sequential_nonlinear_problems():
     bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
 
     # Use SNES options for the nonlinear solver
+    direct_options = {
+        "ksp_monitor": None,
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "pc_factor_mat_solver_type": "mumps",
+    }
     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",
     }
+    options.update(direct_options)
 
     # 5. Solve the Cascade
-    problem1 = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    problem1 = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=direct_options)
     problem1.solve()
 
-    problem2 = NonlinearProblem(F2, u=u2, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    problem2 = NonlinearProblem(F2, u=u2, bcs=[bc], petsc_options=options, adjoint_petsc_options=direct_options)
     problem2.solve()
 
     # 6. Objective (using the cubed error to ensure a 3.0 Hessian rate)
diff --git a/tests/test_solver_reuse.py b/tests/test_solver_reuse.py
index 558e1d1..15eaffa 100644
--- a/tests/test_solver_reuse.py
+++ b/tests/test_solver_reuse.py
@@ -199,13 +199,17 @@ def test_nonlinear_recompute_does_not_corrupt_original_control():
     bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))
     bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
 
-    options = {
-        "snes_error_if_not_converged": True,
+    direct_options = {
+        "ksp_monitor": None,
         "ksp_type": "preonly",
         "pc_type": "lu",
         "pc_factor_mat_solver_type": "mumps",
+        "ksp_error_if_not_converged": True,
     }
-    problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    options = {
+        "snes_error_if_not_converged": True,
+    }
+    problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=direct_options)
     problem.solve()
 
     d = pyadjoint.AdjFloat(0.2)
@@ -319,13 +323,18 @@ def test_nonlinear_adjoint_lhs_compiled_once():
     bc_val = dolfinx.fem.Constant(mesh, np.dtype(dolfinx.default_scalar_type).type(1.0))
     bc = dolfinx.fem.dirichletbc(bc_val, boundary_dofs, V)
 
-    options = {
-        "snes_error_if_not_converged": True,
+    direct_options = {
+        "ksp_monitor": None,
         "ksp_type": "preonly",
         "pc_type": "lu",
         "pc_factor_mat_solver_type": "mumps",
+        "ksp_error_if_not_converged": True,
     }
-    problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+    options = {
+        "snes_error_if_not_converged": True,
+    }
+    options.update(direct_options)
+    problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=direct_options)
     problem.solve()
 
     d = pyadjoint.AdjFloat(0.2)
@@ -626,16 +635,20 @@ def test_nonlinear_problem_released_by_refcounting_not_gc():
     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)
-
-    options = {
-        "snes_error_if_not_converged": True,
+    direct_options = {
+        "ksp_monitor": None,
         "ksp_type": "preonly",
         "pc_type": "lu",
         "pc_factor_mat_solver_type": "mumps",
+        "ksp_error_if_not_converged": True,
+    }
+    options = {
+        "snes_error_if_not_converged": True,
     }
+    options.update(direct_options)
     gc.disable()
     try:
-        problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=options)
+        problem = NonlinearProblem(F1, u=u1, bcs=[bc], petsc_options=options, adjoint_petsc_options=direct_options)
         problem.solve()
 
         problem_ref = weakref.ref(problem)
@@ -732,23 +745,21 @@ def test_nonlinear_problem_rebuilt_after_garbage_collection():
 
     options = {
         "snes_error_if_not_converged": True,
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "pc_factor_mat_solver_type": "mumps",
     }
-    adjoint_options = {
+    direct_options = {
         "ksp_type": "preonly",
         "pc_type": "lu",
         "ksp_error_if_not_converged": True,
         "pc_factor_mat_solver_type": "mumps",
     }
+    options.update(direct_options)
 
     problem = NonlinearProblem(
         F1,
         u=u1,
         bcs=[bc],
         petsc_options=options,
-        adjoint_petsc_options=adjoint_options,
+        adjoint_petsc_options=direct_options,
         petsc_options_prefix="dxa_nonlinear_rebuild_test_",
     )
     problem.solve()
@@ -847,31 +858,28 @@ def test_nonlinear_blocked_problem_templates_compiled_once():
     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)
-
+    ksp_options = {
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+    }
     forward_options = {
         "snes_type": "newtonls",
         "snes_error_if_not_converged": True,
         "snes_atol": 1e-9,
         "snes_rtol": 1e-9,
         "snes_stol": 1e-12,
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "pc_factor_mat_solver_type": "mumps",
-    }
-    adjoint_options = {
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "ksp_error_if_not_converged": True,
-        "pc_factor_mat_solver_type": "mumps",
     }
+    forward_options.update(ksp_options)
     problem = NonlinearProblem(
         [F0, F1],
         u=[uh, ph],
         bcs=[bc],
         petsc_options_prefix="dxa_blocked_nonlinear_reuse_test_",
         petsc_options=forward_options,
-        adjoint_petsc_options=adjoint_options,
-        tlm_petsc_options=adjoint_options,
+        adjoint_petsc_options=ksp_options,
+        tlm_petsc_options=ksp_options,
     )
     problem.solve()
 
@@ -914,3 +922,80 @@ def test_nonlinear_blocked_problem_templates_compiled_once():
     assert problem._get_or_build_hessian_templates() is hessian_templates, (
         "blocked Hessian templates were rebuilt after evaluating at a new point"
     )
+
+
+def test_recompute_checkpoint_and_tlm_value_preserve_overloaded_function_type():
+    """A block's recomputed checkpoint and TLM value must stay the overloaded
+    ``dolfinx_adjoint`` ``Function``, never silently downgrade to the plain
+    ``dolfinx.fem.Function`` a bare ``.copy()`` would produce.
+
+    Regression test for ``_ProblemBlockBase.recompute_component``/
+    ``LinearProblemBlock.__init__`` (and the ``NonlinearProblemBlock`` equivalent)
+    switching from ``.copy()`` -- which is hardcoded, in ``dolfinx.fem.Function``, to
+    always return the plain base class regardless of the source's concrete type -- to
+    ``Function._ad_create_checkpoint()``, which correctly preserves it.
+    """
+    pyadjoint.get_working_tape().clear_tape()
+
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4)
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+
+    uh = Function(V, name="state")
+    v = ufl.TestFunction(V)
+    u_trial = ufl.TrialFunction(V)
+    m = Function(V, name="control")
+    m.interpolate(lambda x: 1.0 + x[0] ** 2)
+
+    a = m * ufl.inner(ufl.grad(u_trial), ufl.grad(v)) * ufl.dx
+    L = ufl.inner(dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0)), 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)
+
+    petsc_options = {
+        "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)
+    problem.solve()
+
+    # create_block_variable() (called by problem.solve()'s own tape recording)
+    # stores the new BlockVariable back onto uh itself, so uh.block_variable is
+    # exactly the one LinearProblemBlock will recompute into below.
+    output_bv = uh.block_variable
+
+    # Using uh as an AssembleBlock coefficient here is what first turns its
+    # block_variable into a dependency of another block, which is what triggers
+    # pyadjoint to freeze a real checkpoint for it (BlockVariable.will_add_as_dependency
+    # -> save_output -> Function._ad_create_checkpoint) -- exercised by every ordinary
+    # J = assemble_scalar(f(uh) * dx) call, not something special-cased for this test.
+    J = assemble_scalar(uh * uh * ufl.dx)
+    control = pyadjoint.Control(m)
+    Jh = pyadjoint.ReducedFunctional(J, control)
+
+    assert isinstance(output_bv.checkpoint, Function), (
+        "the checkpoint frozen when uh was first used as AssembleBlock's dependency "
+        "should already be the overloaded Function"
+    )
+
+    # Re-evaluating at a new point forces LinearProblemBlock.prepare_recompute_component/
+    # recompute_component; the checkpoint above already exists, so this exercises the
+    # in-place-reuse branch (mirroring InterpolationBlock/Firedrake's GenericSolveBlock).
+    m2 = Function(V)
+    m2.interpolate(lambda x: 2.0 + np.sin(x[0]))
+    Jh(m2)
+    assert isinstance(output_bv.saved_output, Function), (
+        "recompute_component's checkpoint reuse must preserve the overloaded Function type"
+    )
+
+    # A Hessian evaluation drives a TLM sweep first, exercising evaluate_tlm_component.
+    dm = Function(V)
+    dm.interpolate(lambda x: np.cos(np.pi * x[0]))
+    Jh.hessian(dm)
+    assert isinstance(output_bv.tlm_value, Function), (
+        "evaluate_tlm_component's returned tlm_value must be the overloaded Function type"
+    )
diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index f4b1646..17f81ad 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -177,10 +177,8 @@ def _navier_stokes(mesh):
         "snes_atol": 1e-9,
         "snes_rtol": 1e-9,
         "snes_stol": 1e-12,
-        "ksp_type": "preonly",
-        "pc_type": "lu",
-        "pc_factor_mat_solver_type": "mumps",
     }
+    forward_options.update(direct_solve)
     problem = NonlinearProblem(
         [F0, F1],
         u=[uh, ph],

From 54f7ec469f646a87ea3f927acd0c42e59e4bc51c Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 07:23:06 +0000
Subject: [PATCH 05/22] Fix minor check

---
 src/dolfinx_adjoint/types/dirichletbc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dolfinx_adjoint/types/dirichletbc.py b/src/dolfinx_adjoint/types/dirichletbc.py
index f17e6e3..548f19f 100644
--- a/src/dolfinx_adjoint/types/dirichletbc.py
+++ b/src/dolfinx_adjoint/types/dirichletbc.py
@@ -50,7 +50,7 @@ def __init__(self, g: Function, dofs: npt.NDArray[np.int32], **kwargs):
         bc_kwargs: dict[str, Any] = {}
         # If dolfinx-version is 0.12 we need to pass the following
         # due to https://github.com/FEniCS/dolfinx/pull/4342/
-        if Version(dolfinx.__version__).minor >= 11:
+        if Version(dolfinx.__version__).minor > 11:
             bc_kwargs["V"] = g.function_space
             bc_kwargs["g"] = g
 

From e4a83f7cfefaf464d4b2651bd322068d1fccc4c3 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 08:39:37 +0000
Subject: [PATCH 06/22] Update options

---
 tests/test_blocked_problem.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index 97008ad..c592e7c 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -185,9 +185,9 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     forward_options = {
         "snes_type": "newtonls",
         "snes_error_if_not_converged": True,
-        "snes_atol": 1e-9,
-        "snes_rtol": 1e-9,
-        "snes_stol": 1e-12,
+        "snes_atol": 1e-8,
+        "snes_rtol": 1e-8,
+        "snes_monitor": None,
     }
     forward_options.update(direct_solve)
     problem = NonlinearProblem(

From 43c62b3801ae7bb48745d65dbfcb774be8a012fb Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 08:48:24 +0000
Subject: [PATCH 07/22] Further tolerance fixes

---
 tests/test_tlm_update.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index 17f81ad..a5dc923 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -174,9 +174,8 @@ def _navier_stokes(mesh):
         # step can make backtracking line search report DIVERGED_LINE_SEARCH.  An
         # explicit absolute tolerance lets it recognize "already converged" and
         # exit immediately instead.
-        "snes_atol": 1e-9,
-        "snes_rtol": 1e-9,
-        "snes_stol": 1e-12,
+        "snes_atol": 1e-8,
+        "snes_rtol": 1e-8,
     }
     forward_options.update(direct_solve)
     problem = NonlinearProblem(

From c2c36b63e75d751c16ea6d2e75c3629602984859 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 09:23:57 +0000
Subject: [PATCH 08/22] Sets are unsorted. Thus the coefficient set might be
 ordered differently on each process. The placeholder creation then deadlocked
 as Function initialization is a collective operation. Add ufl_id to Function
 and Constant to have something to sort them by.

Another improvement is to use direct solvers in the test (not relying on petsc defaults) and setting adjoint and tlm options as well
---
 src/dolfinx_adjoint/solvers.py        |  5 +--
 src/dolfinx_adjoint/types/function.py |  8 +++++
 tests/test_hessian.py                 | 51 +++++++++++++++++++++++++--
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/src/dolfinx_adjoint/solvers.py b/src/dolfinx_adjoint/solvers.py
index 75d28f6..5882c37 100644
--- a/src/dolfinx_adjoint/solvers.py
+++ b/src/dolfinx_adjoint/solvers.py
@@ -791,9 +791,10 @@ def __init__(
             coefficients |= collect_coefficients(P)
             if set(u_list).issubset(coefficients):
                 raise ValueError("The unknown `u` should not be part of the coefficients of a linear problem.")
-
+        # Has to be sorted when creating placeholders, as function creation is a collective operation
+        sorted_coefficients = sorted(coefficients, key=lambda c: c.ufl_id())
         self._value_placeholders: dict[dolfinx.fem.Function, dolfinx.fem.Function] = {
-            c: dolfinx.fem.Function(c.function_space) for c in coefficients
+            c: dolfinx.fem.Function(c.function_space) for c in sorted_coefficients
         }
         a_R, L_R, P_R = recursive_replace((a, L, P), self._value_placeholders)  # type: ignore[misc]
         super().__init__(
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index eae65a7..a55103e 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -13,6 +13,7 @@
     register_overloaded_type,
 )
 from pyadjoint.tape import no_annotations
+from ufl.core.ufl_id import attach_ufl_id
 
 from ..blocks._vector import _SpecialVector, _vector
 from ..blocks.assembly import assemble_compiled_form
@@ -35,6 +36,7 @@ def _create_function(
     return Function(V, x=x, annotate=False)
 
 
+@attach_ufl_id
 class Function(dolfinx.fem.Function, FloatingType):
     """A class overloading `dolfinx.fem.Function` to support it being used as a control variable
     in the adjoint framework.
@@ -56,6 +58,8 @@ def __init__(
         dtype: npt.DTypeLike = dolfinx.default_scalar_type,
         **kwargs,
     ):
+        ufl_id = kwargs.pop("ufl_id", None)
+        self._ufl_id = self._init_ufl_id(ufl_id)
         super(Function, self).__init__(
             V,
             x,
@@ -249,6 +253,7 @@ def x(self) -> dolfinx.la.Vector:
         return self._x
 
 
+@attach_ufl_id
 class Constant(Function):
     """A class overloading {py:class}`dolfinx.fem.Constant`
     to support it being used as a control variable in
@@ -274,7 +279,10 @@ def __init__(
         self,
         domain: dolfinx.mesh.Mesh,
         c: float | numpy.floating | complex | numpy.complexfloating | typing.Sequence | numpy.ndarray,
+        ufl_id: int | None = None,
     ):
+        self._ufl_id = self._init_ufl_id(ufl_id)
+
         value_shape = numpy.shape(c)
         try:
             el = basix.ufl.real_element(domain.basix_cell(), value_shape=numpy.shape(c))
diff --git a/tests/test_hessian.py b/tests/test_hessian.py
index a8a05c7..19a22d8 100644
--- a/tests/test_hessian.py
+++ b/tests/test_hessian.py
@@ -40,8 +40,23 @@ def test_constant_hessian():
     bc = dolfinx_adjoint.dirichletbc(uD, boundary_dofs)
 
     # Solve and tape the PDE
+    ksp_options = {
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+        "ksp_monitor": None,
+    }
     u_sol = dolfinx_adjoint.Function(V, name="State")
-    problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol)
+    problem = dolfinx_adjoint.LinearProblem(
+        a,
+        L,
+        bcs=[bc],
+        u=u_sol,
+        petsc_options=ksp_options,
+        adjoint_petsc_options=ksp_options,
+        tlm_petsc_options=ksp_options,
+    )
     problem.solve()
 
     # ==========================================
@@ -129,7 +144,22 @@ def test_constant_hessian_linear_source():
     bc = dolfinx_adjoint.dirichletbc(u_bc, boundary_dofs)
 
     u_sol = dolfinx_adjoint.Function(V)
-    problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol)
+    ksp_options = {
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+        "ksp_monitor": None,
+    }
+    problem = dolfinx_adjoint.LinearProblem(
+        a,
+        L,
+        bcs=[bc],
+        u=u_sol,
+        petsc_options=ksp_options,
+        adjoint_petsc_options=ksp_options,
+        tlm_petsc_options=ksp_options,
+    )
     problem.solve()
 
     J_form = 0.5 * ufl.inner(u_sol, u_sol) * ufl.dx
@@ -172,7 +202,22 @@ def test_constant_hessian_linear_operator():
     bc = dolfinx_adjoint.dirichletbc(u_bc, boundary_dofs)
 
     u_sol = dolfinx_adjoint.Function(V)
-    problem = dolfinx_adjoint.LinearProblem(a, L, bcs=[bc], u=u_sol)
+    ksp_options = {
+        "ksp_type": "preonly",
+        "pc_type": "lu",
+        "ksp_error_if_not_converged": True,
+        "pc_factor_mat_solver_type": "mumps",
+        "ksp_monitor": None,
+    }
+    problem = dolfinx_adjoint.LinearProblem(
+        a,
+        L,
+        bcs=[bc],
+        u=u_sol,
+        petsc_options=ksp_options,
+        adjoint_petsc_options=ksp_options,
+        tlm_petsc_options=ksp_options,
+    )
     problem.solve()
 
     J_form = 0.5 * ufl.inner(u_sol, u_sol) * ufl.dx

From fd48e288e8110527f5751d627e92ea1d85a5407e Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 10:50:59 +0000
Subject: [PATCH 09/22] Tweak step sizes n taylor test and add scatter forward
 in adjoint/tlm problem

---
 src/dolfinx_adjoint/petsc_utils.py |  2 +-
 tests/test_blocked_problem.py      | 17 ++++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/dolfinx_adjoint/petsc_utils.py b/src/dolfinx_adjoint/petsc_utils.py
index 3a27e50..c295f8b 100644
--- a/src/dolfinx_adjoint/petsc_utils.py
+++ b/src/dolfinx_adjoint/petsc_utils.py
@@ -99,7 +99,7 @@ def solve(
             except RuntimeError:
                 bcs0 = dolfinx.fem.bcs.bcs_by_block(dolfinx.fem.forms.extract_spaces(self._L), self.bcs)  # type: ignore
                 dolfinx.fem.petsc.set_bc(self._b, bcs0, alpha=0.0)
-
+        self._b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)  # type: ignore
         # Solve linear system and update ghost values in the solution
         self._solver.solve(self._b, self._x)
         dolfinx.la.petsc._ghost_update(self._x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore
diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index c592e7c..23eb0c4 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -20,7 +20,7 @@
 
 @pytest.fixture(scope="module")
 def mesh_2D():
-    return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 7)
+    return dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 27, 32)
 
 
 @pytest.mark.parametrize("use_mixed_space", [True, False])
@@ -50,7 +50,7 @@ def L1(mesh, q):
 
     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]))
+    f.interpolate(lambda x: (np.sin(x[0]), -2 * x[1]))
 
     if use_mixed_space:
         W = ufl.MixedFunctionSpace(*[V, Q])
@@ -102,7 +102,7 @@ def L1(mesh, q):
     d.interpolate(lambda x: (10 * x[0], x[1]))
 
     e = Function(Z)
-    e.interpolate(lambda x: (1e3 * np.sin(x[1]), 1e3 * x[0]))
+    e.interpolate(lambda x: (np.sin(x[1]), 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}"
 
@@ -110,7 +110,10 @@ def L1(mesh, q):
     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}"
 
+    # Scale perturbation for hessian
     Jh(d)
+    e.x.array[:] *= 200
+    e.x.scatter_forward()
     dJdm = Jh.derivative()._ad_dot(e)
     hessian = Jh.hessian(e)
     dHddu = hessian._ad_dot(e)
@@ -118,9 +121,11 @@ def L1(mesh, q):
     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)))
+    z.interpolate(lambda x: (3 * np.sin(x[1]), -5 * (x[0] ** 2)))
     f = Function(Z)
-    f.interpolate(lambda x: (1e4 * x[0], 1e5 * np.sin(x[1])))
+    f.interpolate(lambda x: (1e3 * x[0] ** 2, 1e2 * np.sin(x[1])))
+    f.x.array[:] *= 10_030
+    f.x.scatter_forward()
     Jh(z)
     dJdm = Jh.derivative()._ad_dot(f)
     hessian = Jh.hessian(f)
@@ -223,6 +228,7 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     assert np.isclose(min_rate, 2.0, rtol=1e-1, atol=1e-1), f"Expected convergence rate close to 2.0, got {min_rate}"
 
     Jh(d)
+    e.x.array[:] *= 300
     dJdm = Jh.derivative()._ad_dot(e)
     hessian = Jh.hessian(e)
     dHddu = hessian._ad_dot(e)
@@ -237,6 +243,7 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     h2 = Function(Z)
     h2.interpolate(lambda x: 0.5 * np.cos(4 * x[0]))
     Jh(mu2)
+    h2.x.array[:] *= 103
     dJdm = Jh.derivative()._ad_dot(h2)
     hessian = Jh.hessian(h2)
     dHddu = hessian._ad_dot(h2)

From 74916d3730cf130be420e1f8d9cb97a5f09dcc68 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 11:15:41 +0000
Subject: [PATCH 10/22] Make sure perturbations are physical

---
 tests/test_blocked_problem.py | 148 ++++++++++++++++++----------------
 1 file changed, 79 insertions(+), 69 deletions(-)

diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index 23eb0c4..2115e0c 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -98,40 +98,45 @@ def L1(mesh, q):
 
     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: (np.sin(x[1]), 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}"
-
-    # Scale perturbation for hessian
-    Jh(d)
-    e.x.array[:] *= 200
-    e.x.scatter_forward()
-    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: (3 * np.sin(x[1]), -5 * (x[0] ** 2)))
-    f = Function(Z)
-    f.interpolate(lambda x: (1e3 * x[0] ** 2, 1e2 * np.sin(x[1])))
-    f.x.array[:] *= 10_030
-    f.x.scatter_forward()
-    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}"
+    with pyadjoint.stop_annotating():
+        d = Function(Z)
+        d.interpolate(lambda x: (10 * x[0], x[1]))
+
+        e = Function(Z)
+        e.interpolate(lambda x: (np.sin(x[1]), x[0] ** 2))
+        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}"
+        )
+
+        # Scale perturbation for hessian
+        Jh(d)
+        e.x.array[:] *= 201
+        e.x.scatter_forward()
+        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: (3 * np.sin(x[1]) + x[1] * x[0], -5 * (x[0]) + (1 - x[1])))
+        f = Function(Z)
+        f.interpolate(lambda x: (x[1] ** 2, x[0] ** 2))  # NOTE: Has to be divergence free
+        f.x.array[:] *= 102
+        f.x.scatter_forward()
+        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}"
 
 
 @pytest.mark.parametrize("use_mixed_space", [True, False])
@@ -151,7 +156,7 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     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]))
+    mu.interpolate(lambda x: 3.0 + 0.5 * np.sin(np.pi * x[0]))
 
     uh, ph = Function(V, name="velocity"), Function(Q, name="pressure")
 
@@ -215,37 +220,42 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
 
     control = pyadjoint.Control(mu)
     Jh = pyadjoint.ReducedFunctional(J, control)
-    d = Function(Z)
-    d.interpolate(lambda x: 1.0 + 0.3 * np.cos(np.pi * x[1]))
-    e = Function(Z)
-    e.interpolate(lambda x: 0.2 * np.sin(3 * 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-1, atol=1e-1), f"Expected convergence rate close to 2.0, got {min_rate}"
-
-    Jh(d)
-    e.x.array[:] *= 300
-    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}"
-
-    # A second, independent evaluation point/direction: a cached-but-unrefreshed
-    # adjoint/TLM/Hessian operator (see tests/test_tlm_update.py) could pass the
-    # check above yet still be silently wrong here.
-    mu2 = Function(Z)
-    mu2.interpolate(lambda x: 2.0 + np.sin(x[1]))
-    h2 = Function(Z)
-    h2.interpolate(lambda x: 0.5 * np.cos(4 * x[0]))
-    Jh(mu2)
-    h2.x.array[:] *= 103
-    dJdm = Jh.derivative()._ad_dot(h2)
-    hessian = Jh.hessian(h2)
-    dHddu = hessian._ad_dot(h2)
-    min_rate = pyadjoint.taylor_test(Jh, mu2, h2, 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}"
+    with pyadjoint.stop_annotating():
+        d = Function(Z)
+        d.interpolate(lambda x: 25.0 + 0.3 * np.cos(np.pi * x[1]))
+        e = Function(Z)
+        e.interpolate(lambda x: 0.1 * np.sin(3 * x[0]))
+        e.x.array[:] *= 20  # NOTE: min(d)-max(e) > 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-1, atol=1e-1), (
+            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}"
+
+        # A second, independent evaluation point/direction: a cached-but-unrefreshed
+        # adjoint/TLM/Hessian operator (see tests/test_tlm_update.py) could pass the
+        # check above yet still be silently wrong here.
+        mu2 = Function(Z)
+        mu2.interpolate(lambda x: 20.0 + np.sin(x[1]))
+        h2 = Function(Z)
+        h2.interpolate(lambda x: 0.5 * np.cos(4 * x[0]))  # NOTE: min(mu2)-max(h2) > 0
+        h2.x.array[:] *= 1e1
+        Jh(mu2)
+        dJdm = Jh.derivative()._ad_dot(h2)
+        hessian = Jh.hessian(h2)
+        dHddu = hessian._ad_dot(h2)
+        min_rate = pyadjoint.taylor_test(Jh, mu2, h2, 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}"

From f3a34afc046ce50deceb787c4c3374778386c2bf Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 11:22:22 +0000
Subject: [PATCH 11/22] More ensuring of physical quantities

---
 tests/test_tlm_update.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index a5dc923..28fd9cd 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -51,7 +51,7 @@ def _viscous_stokes(mesh):
     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]))
+    mu.interpolate(lambda x: 15.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)
@@ -104,8 +104,8 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another
     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]))
+    m1.interpolate(lambda x: 15.0 + 0.5 * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 25.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]))
 
@@ -137,7 +137,7 @@ def _navier_stokes(mesh):
     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]))
+    mu.interpolate(lambda x: 25.0 + 0.5 * np.sin(np.pi * x[0]))
 
     uh, ph = Function(V, name="velocity"), Function(Q, name="pressure")
     v, q = ufl.TestFunction(V), ufl.TestFunction(Q)
@@ -211,7 +211,8 @@ def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh
     Jh, Z = _navier_stokes(mesh_2D)
 
     m2 = Function(Z)
-    m2.interpolate(lambda x: 2.0 + 0.5 * np.cos(np.pi * x[1]))
+    m2.interpolate(lambda x: 25.0 + 0.5 * np.cos(np.pi * x[1]))
+    # Ensure min(m2) - max(h) > 0, so that the Taylor test's finite-difference perturbation doesn't
     h = Function(Z)
     h.interpolate(lambda x: 1.0 + 0.3 * np.sin(3 * x[0]))
 
@@ -234,7 +235,7 @@ def _diffusive_poisson(mesh):
     dx = ufl.Measure("dx", domain=mesh)
 
     m = Function(Z, name="diffusivity")
-    m.interpolate(lambda x: 1.0 + 0.5 * np.sin(np.pi * x[0]))
+    m.interpolate(lambda x: 23.0 + 0.5 * np.sin(np.pi * x[0]))
 
     u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
     x = ufl.SpatialCoordinate(mesh)
@@ -281,8 +282,8 @@ def test_hessian_is_independent_of_previous_evaluation_points_scalar(warm_up_at_
     Jh, Z = _diffusive_poisson(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]))
+    m1.interpolate(lambda x: 25.0 + 0.5 * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 40.0 + 0.5 * np.cos(np.pi * x[1]))
     h = Function(Z)
     h.interpolate(lambda x: 1.0 + 0.3 * np.sin(3 * x[0]))
 
@@ -304,8 +305,8 @@ def test_hessian_mpi_breakdown(mesh_2D):
     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]))
+    m1.interpolate(lambda x: 17.0 + 0.5 * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 13.2 + 0.2 * np.cos(np.pi * x[1]))
     h = Function(Z)
     h.interpolate(lambda x: 3.0 + 2.0 * np.sin(3 * x[0]))
 

From 746e555f0e4febcaef3e308e049758f10c193582 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 11:30:26 +0000
Subject: [PATCH 12/22] Try stricter tolerance

---
 tests/test_blocked_problem.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index 2115e0c..377c309 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -195,8 +195,8 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     forward_options = {
         "snes_type": "newtonls",
         "snes_error_if_not_converged": True,
-        "snes_atol": 1e-8,
-        "snes_rtol": 1e-8,
+        "snes_atol": 1e-12,
+        "snes_rtol": 1e-12,
         "snes_monitor": None,
     }
     forward_options.update(direct_solve)

From ce36b85f1a042f67476d5d1eefda75198132afcc Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 13:02:02 +0000
Subject: [PATCH 13/22] Try again

---
 tests/test_blocked_problem.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index 377c309..4e3498f 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -195,8 +195,8 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     forward_options = {
         "snes_type": "newtonls",
         "snes_error_if_not_converged": True,
-        "snes_atol": 1e-12,
-        "snes_rtol": 1e-12,
+        "snes_atol": 1e-9,
+        "snes_rtol": 1e-9,
         "snes_monitor": None,
     }
     forward_options.update(direct_solve)
@@ -220,13 +220,13 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
 
     control = pyadjoint.Control(mu)
     Jh = pyadjoint.ReducedFunctional(J, control)
+    baseline = 10
     with pyadjoint.stop_annotating():
         d = Function(Z)
-        d.interpolate(lambda x: 25.0 + 0.3 * np.cos(np.pi * x[1]))
+        d.interpolate(lambda x: 2 * baseline + 0.3 * baseline * np.cos(np.pi * x[1]))
         e = Function(Z)
-        e.interpolate(lambda x: 0.1 * np.sin(3 * x[0]))
-        e.x.array[:] *= 20  # NOTE: min(d)-max(e) > 0
-
+        e.interpolate(lambda x: 0.25 * baseline * np.sin(3 * x[0]))
+        assert np.min(d.x.array - e.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
         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}"
@@ -249,10 +249,12 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
         # adjoint/TLM/Hessian operator (see tests/test_tlm_update.py) could pass the
         # check above yet still be silently wrong here.
         mu2 = Function(Z)
-        mu2.interpolate(lambda x: 20.0 + np.sin(x[1]))
+        mu2.interpolate(lambda x: 3 * baseline + 0.8 * baseline * np.sin(x[1]))
         h2 = Function(Z)
-        h2.interpolate(lambda x: 0.5 * np.cos(4 * x[0]))  # NOTE: min(mu2)-max(h2) > 0
-        h2.x.array[:] *= 1e1
+        h2.interpolate(lambda x: 0.2 * baseline + 0.25 * baseline * np.cos(x[0]))  # NOTE: min(mu2)-max(h2) > 0
+        assert np.min(mu2.x.array - h2.x.array) > 0.0, (
+            "Taylor test perturbation must not violate positivity of viscosity"
+        )
         Jh(mu2)
         dJdm = Jh.derivative()._ad_dot(h2)
         hessian = Jh.hessian(h2)

From ef06dbc0f9c48e9c9b600a44df61a7dfa541eccc Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 13:12:16 +0000
Subject: [PATCH 14/22] Similar sanity check in test hessian

---
 tests/test_tlm_update.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index 28fd9cd..f0479bb 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -103,17 +103,20 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _viscous_stokes(mesh_2D)
 
+    baseline = 10
     m1, m2 = Function(Z), Function(Z)
-    m1.interpolate(lambda x: 15.0 + 0.5 * np.sin(np.pi * x[0]))
-    m2.interpolate(lambda x: 25.0 + 0.5 * np.cos(np.pi * x[1]))
+    m1.interpolate(lambda x: 4 * baseline + 0.1 * baseline * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 3 * baseline + 0.5 * baseline * np.cos(np.pi * x[1]))
     h = Function(Z)
-    h.interpolate(lambda x: 10.0 + 3.2 * np.sin(3 * x[0]))
+    h.interpolate(lambda x: baseline + 0.8 * baseline * np.sin(3 * x[0]))
+    assert np.min(m1.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
 
     if warm_up_at_another_point:
         Jh(m1)
         Jh.derivative()
         Jh.hessian(h)
 
+    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
     Jh(m2)
     dJdm = Jh.derivative()._ad_dot(h)
     Hm = Jh.hessian(h)._ad_dot(h)

From 328d67366d62de918273f105e30be30e19389316 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 13:19:41 +0000
Subject: [PATCH 15/22] Reduce baseline

---
 tests/test_tlm_update.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index f0479bb..0458164 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -103,12 +103,12 @@ def test_hessian_is_independent_of_previous_evaluation_points(warm_up_at_another
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _viscous_stokes(mesh_2D)
 
-    baseline = 10
+    baseline = 4
     m1, m2 = Function(Z), Function(Z)
-    m1.interpolate(lambda x: 4 * baseline + 0.1 * baseline * np.sin(np.pi * x[0]))
-    m2.interpolate(lambda x: 3 * baseline + 0.5 * baseline * np.cos(np.pi * x[1]))
+    m1.interpolate(lambda x: 2 * baseline + 0.1 * baseline * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 2 * baseline + 0.5 * baseline * np.cos(np.pi * x[1]))
     h = Function(Z)
-    h.interpolate(lambda x: baseline + 0.8 * baseline * np.sin(3 * x[0]))
+    h.interpolate(lambda x: 0.8 * baseline * np.sin(x[0]))
     assert np.min(m1.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
 
     if warm_up_at_another_point:

From fad3f13df829f7450c00bec2dcf51cd4be0f2c5a Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 13:37:47 +0000
Subject: [PATCH 16/22] Try again

---
 tests/test_tlm_update.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index 0458164..e43c54b 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -12,6 +12,7 @@
 between two evaluations at another value doesn't perturb the second evaluation's result).
 """
 
+
 from mpi4py import MPI
 
 import basix.ufl
@@ -213,12 +214,13 @@ def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _navier_stokes(mesh_2D)
 
+    baseline = 10.0
     m2 = Function(Z)
-    m2.interpolate(lambda x: 25.0 + 0.5 * np.cos(np.pi * x[1]))
+    m2.interpolate(lambda x: baseline + 0.2 * baseline * np.cos(np.pi * x[1]))
     # Ensure min(m2) - max(h) > 0, so that the Taylor test's finite-difference perturbation doesn't
     h = Function(Z)
-    h.interpolate(lambda x: 1.0 + 0.3 * np.sin(3 * x[0]))
-
+    h.interpolate(lambda x: 0.3 * baseline * np.sin(3 * x[0]))
+    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
     Jh(m2)
     dJdm = Jh.derivative()._ad_dot(h)
     Hm = Jh.hessian(h)._ad_dot(h)
@@ -283,18 +285,19 @@ def test_hessian_is_independent_of_previous_evaluation_points_scalar(warm_up_at_
     """
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _diffusive_poisson(mesh_2D)
+    baseline = 10.0
 
     m1, m2 = Function(Z), Function(Z)
-    m1.interpolate(lambda x: 25.0 + 0.5 * np.sin(np.pi * x[0]))
-    m2.interpolate(lambda x: 40.0 + 0.5 * np.cos(np.pi * x[1]))
+    m1.interpolate(lambda x: 2 * baseline + 0.1 * baseline * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 3 * baseline + 0.5 * baseline * np.cos(np.pi * x[1]))
     h = Function(Z)
-    h.interpolate(lambda x: 1.0 + 0.3 * np.sin(3 * x[0]))
-
+    h.interpolate(lambda x: 0.3 * baseline * np.sin(3 * x[0]))
+    assert np.min(m1.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive diffusivity"
+    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive diffusivity"
     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)
@@ -306,13 +309,15 @@ def test_hessian_is_independent_of_previous_evaluation_points_scalar(warm_up_at_
 def test_hessian_mpi_breakdown(mesh_2D):
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _viscous_stokes(mesh_2D)
+    baseline = 23.2
 
     m1, m2 = Function(Z), Function(Z)
-    m1.interpolate(lambda x: 17.0 + 0.5 * np.sin(np.pi * x[0]))
-    m2.interpolate(lambda x: 13.2 + 0.2 * np.cos(np.pi * x[1]))
+    m1.interpolate(lambda x: 2 * baseline + 0.25 * baseline * np.sin(np.pi * x[0]))
+    m2.interpolate(lambda x: 4 * baseline + 0.2 * baseline * np.cos(np.pi * x[1]))
     h = Function(Z)
-    h.interpolate(lambda x: 3.0 + 2.0 * np.sin(3 * x[0]))
-
+    h.interpolate(lambda x: -0.8 * baseline * np.cos(x[1]) * np.sin(3 * x[0]))
+    assert np.min(m1.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive viscosity"
+    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive viscosity"
     # === COLD START ===
     J_cold = float(Jh(m2))
     dJ_cold = Jh.derivative()

From b5a1637ce6be09fdca1abf82c2eb22b61f1cbdc1 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 14:01:18 +0000
Subject: [PATCH 17/22] Try modifying tlm update hesssian navier stokes
 bsaeline

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

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index e43c54b..ce39413 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -12,7 +12,6 @@
 between two evaluations at another value doesn't perturb the second evaluation's result).
 """
 
-
 from mpi4py import MPI
 
 import basix.ufl
@@ -214,7 +213,7 @@ def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh
     pyadjoint.get_working_tape().clear_tape()
     Jh, Z = _navier_stokes(mesh_2D)
 
-    baseline = 10.0
+    baseline = 0.08
     m2 = Function(Z)
     m2.interpolate(lambda x: baseline + 0.2 * baseline * np.cos(np.pi * x[1]))
     # Ensure min(m2) - max(h) > 0, so that the Taylor test's finite-difference perturbation doesn't

From 4d01e128f6284399b1e3d88fa18b0c97ba1efc50 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 14:12:06 +0000
Subject: [PATCH 18/22] More baselines to scale taylor tests

---
 tests/test_blocked_problem.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/tests/test_blocked_problem.py b/tests/test_blocked_problem.py
index 4e3498f..c0ccee7 100644
--- a/tests/test_blocked_problem.py
+++ b/tests/test_blocked_problem.py
@@ -99,11 +99,12 @@ def L1(mesh, q):
     control = pyadjoint.Control(f)
     Jh = pyadjoint.ReducedFunctional(J, control)
     with pyadjoint.stop_annotating():
+        baseline = 100
         d = Function(Z)
-        d.interpolate(lambda x: (10 * x[0], x[1]))
+        d.interpolate(lambda x: (baseline * x[0], 2 * baseline * x[1]))
 
         e = Function(Z)
-        e.interpolate(lambda x: (np.sin(x[1]), x[0] ** 2))
+        e.interpolate(lambda x: (10 * baseline * np.sin(x[1]), 3 * baseline * x[0] ** 2))
         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}"
@@ -117,8 +118,6 @@ def L1(mesh, q):
 
         # Scale perturbation for hessian
         Jh(d)
-        e.x.array[:] *= 201
-        e.x.scatter_forward()
         dJdm = Jh.derivative()._ad_dot(e)
         hessian = Jh.hessian(e)
         dHddu = hessian._ad_dot(e)
@@ -126,10 +125,13 @@ def L1(mesh, q):
         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: (3 * np.sin(x[1]) + x[1] * x[0], -5 * (x[0]) + (1 - x[1])))
+        z.interpolate(
+            lambda x: (5 * baseline * np.sin(x[1]) + 0.3 * baseline * x[1] * x[0], -2 * baseline * (x[0]) + (1 - x[1]))
+        )
         f = Function(Z)
-        f.interpolate(lambda x: (x[1] ** 2, x[0] ** 2))  # NOTE: Has to be divergence free
-        f.x.array[:] *= 102
+        f.interpolate(
+            lambda x: (0.8 * baseline * x[1] ** 2, 2 * baseline * x[0] ** 2)
+        )  # NOTE: Has to be divergence free
         f.x.scatter_forward()
         Jh(z)
         dJdm = Jh.derivative()._ad_dot(f)
@@ -156,7 +158,8 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
     dx = ufl.Measure("dx", domain=mesh)
 
     mu = Function(Z, name="viscosity")
-    mu.interpolate(lambda x: 3.0 + 0.5 * np.sin(np.pi * x[0]))
+    baseline = 0.08
+    mu.interpolate(lambda x: baseline + 0.5 * baseline * np.sin(np.pi * x[0]))
 
     uh, ph = Function(V, name="velocity"), Function(Q, name="pressure")
 
@@ -220,12 +223,12 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
 
     control = pyadjoint.Control(mu)
     Jh = pyadjoint.ReducedFunctional(J, control)
-    baseline = 10
+    baseline = 0.08
     with pyadjoint.stop_annotating():
         d = Function(Z)
         d.interpolate(lambda x: 2 * baseline + 0.3 * baseline * np.cos(np.pi * x[1]))
         e = Function(Z)
-        e.interpolate(lambda x: 0.25 * baseline * np.sin(3 * x[0]))
+        e.interpolate(lambda x: 0.5 * baseline * np.sin(3 * x[0]))
         assert np.min(d.x.array - e.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
         min_rate = pyadjoint.taylor_test(Jh, d, e, dJdm=0)
         assert np.isclose(min_rate, 1.0, rtol=1e-1, atol=1e-1), (
@@ -251,7 +254,7 @@ def test_nonlinear_solver(use_mixed_space: bool, mesh_2D):
         mu2 = Function(Z)
         mu2.interpolate(lambda x: 3 * baseline + 0.8 * baseline * np.sin(x[1]))
         h2 = Function(Z)
-        h2.interpolate(lambda x: 0.2 * baseline + 0.25 * baseline * np.cos(x[0]))  # NOTE: min(mu2)-max(h2) > 0
+        h2.interpolate(lambda x: 0.4 * baseline + 0.9 * baseline * np.cos(x[0]))  # NOTE: min(mu2)-max(h2) > 0
         assert np.min(mu2.x.array - h2.x.array) > 0.0, (
             "Taylor test perturbation must not violate positivity of viscosity"
         )

From caa0536388fbaf91146d8f751ec907fbfbd26090 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 14:34:21 +0000
Subject: [PATCH 19/22] Is this the last update?

---
 tests/test_tlm_update.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index ce39413..a304356 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -218,8 +218,8 @@ def test_hessian_is_independent_of_previous_evaluation_points_navier_stokes(mesh
     m2.interpolate(lambda x: baseline + 0.2 * baseline * np.cos(np.pi * x[1]))
     # Ensure min(m2) - max(h) > 0, so that the Taylor test's finite-difference perturbation doesn't
     h = Function(Z)
-    h.interpolate(lambda x: 0.3 * baseline * np.sin(3 * x[0]))
-    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positivity of viscosity"
+    h.interpolate(lambda x: 0.4 * baseline * np.sin(3 * x[0]))
+    assert np.min(m2.x.array - h.x.array) > 0.01, "Taylor test perturbation must not violate positivity of viscosity"
     Jh(m2)
     dJdm = Jh.derivative()._ad_dot(h)
     Hm = Jh.hessian(h)._ad_dot(h)

From 27606c040f384cbd9bcf2215290ba481a8846deb Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 14:43:23 +0000
Subject: [PATCH 20/22] Add monitor for debug

---
 tests/test_tlm_update.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index a304356..87d4ce6 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -29,6 +29,7 @@
     "pc_type": "lu",
     "ksp_error_if_not_converged": True,
     "pc_factor_mat_solver_type": "mumps",
+    "ksp_monitor": None,
 }
 
 
@@ -177,6 +178,7 @@ def _navier_stokes(mesh):
         # step can make backtracking line search report DIVERGED_LINE_SEARCH.  An
         # explicit absolute tolerance lets it recognize "already converged" and
         # exit immediately instead.
+        "snes_monitor": None,
         "snes_atol": 1e-8,
         "snes_rtol": 1e-8,
     }
@@ -315,8 +317,8 @@ def test_hessian_mpi_breakdown(mesh_2D):
     m2.interpolate(lambda x: 4 * baseline + 0.2 * baseline * np.cos(np.pi * x[1]))
     h = Function(Z)
     h.interpolate(lambda x: -0.8 * baseline * np.cos(x[1]) * np.sin(3 * x[0]))
-    assert np.min(m1.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive viscosity"
-    assert np.min(m2.x.array - h.x.array) > 0.0, "Taylor test perturbation must not violate positive viscosity"
+    assert np.min(m1.x.array - h.x.array) > 0.05, "Taylor test perturbation must not violate positive viscosity"
+    assert np.min(m2.x.array - h.x.array) > 0.05, "Taylor test perturbation must not violate positive viscosity"
     # === COLD START ===
     J_cold = float(Jh(m2))
     dJ_cold = Jh.derivative()

From 5a324ed6d9ed284107ab738fa475af811a138fc3 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 14:55:15 +0000
Subject: [PATCH 21/22] Try null pivot detection and reordering

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

diff --git a/tests/test_tlm_update.py b/tests/test_tlm_update.py
index 87d4ce6..ab5fc08 100644
--- a/tests/test_tlm_update.py
+++ b/tests/test_tlm_update.py
@@ -30,6 +30,8 @@
     "ksp_error_if_not_converged": True,
     "pc_factor_mat_solver_type": "mumps",
     "ksp_monitor": None,
+    "mat_mumps_icntl_24": 1,
+    "pc_factor_mat_ordering_type": "rcm",
 }
 
 

From 469d1641fc5bf0970a32a5d3b774ef982e72cea8 Mon Sep 17 00:00:00 2001
From: Joergen Schartum Dokken <dokken@simula.no>
Date: Wed, 2 Sep 2026 15:00:43 +0000
Subject: [PATCH 22/22] Add reorder and pivot detection

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

diff --git a/tests/test_nonlinear_problem.py b/tests/test_nonlinear_problem.py
index 9cf1528..d6ab511 100644
--- a/tests/test_nonlinear_problem.py
+++ b/tests/test_nonlinear_problem.py
@@ -49,6 +49,8 @@ def test_sequential_nonlinear_problems():
         "ksp_type": "preonly",
         "pc_type": "lu",
         "pc_factor_mat_solver_type": "mumps",
+        "mat_mumps_icntl_24": 1,
+        "pc_factor_mat_ordering_type": "rcm",
     }
     options = {
         "snes_monitor": None,