From ebe2cce22a9434f52e457e7de0fcca47ce96b7b1 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Fri, 21 Aug 2026 13:31:00 +0000
Subject: [PATCH 1/5] Fix unsafe id()-keyed matrix cache and stale recompute
 output in InterpolationBlock

The interpolation matrix cache was a module-level dict keyed by
(id(space_from), id(space_to)) without holding strong references, so a
garbage-collected FunctionSpace's id could be reused by an unrelated
object and silently return the wrong matrix. Cache the matrix as a
per-instance attribute on InterpolationBlock instead, tied to the
block's own space references.

recompute_component also wrapped a separately cached Function instead
of updating the tape's real output object, permanently disconnecting
it from the Python object returned by interpolate(). It now mutates
block_variable.saved_output in place, matching FunctionAssignBlock's
convention. Note this only holds until the output is used as a
dependency elsewhere, at which point pyadjoint freezes its own
checkpoint copy (an existing, codebase-wide characteristic, not
specific to interpolation) - documented inline.

Also fixes the mypy failures in this file: the missing
`dolfinx.fem.petsc` import that `petsc_mat=True` silently relied on
via import order, and the monkey-patched `_row_vec`/`_col_vec`
attributes on MatrixCSR.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 src/dolfinx_adjoint/blocks/interpolation.py | 109 +++++++++++---------
 1 file changed, 61 insertions(+), 48 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index 8cf8b58..e26a5f2 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -4,24 +4,24 @@
 from typing import Callable
 
 import dolfinx
+import dolfinx.fem.petsc
 from pyadjoint import Block
-from pyadjoint.overloaded_type import create_overloaded_object
 
 if typing.TYPE_CHECKING:
     from petsc4py import PETSc
 
-# Global cache to prevent redundant matrix assembly across multiple blocks/iterations
-_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int], dolfinx.la.MatrixCSR] = {}
-
 
 def attach_working_array(mat: dolfinx.la.MatrixCSR):
     """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication."""
-    if not hasattr(mat, "_row_vec"):
-        mat._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
-    if not hasattr(mat, "_col_vec"):
-        mat._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
-    mat._row_vec.array[:] = 0.0
-    mat._col_vec.array[:] = 0.0
+    # mat._row_vec/_col_vec are monkey-patched on; cast to Any so mypy doesn't
+    # flag the dynamic attributes.
+    m = typing.cast(typing.Any, mat)
+    if not hasattr(m, "_row_vec"):
+        m._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
+    if not hasattr(m, "_col_vec"):
+        m._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
+    m._row_vec.array[:] = 0.0
+    m._col_vec.array[:] = 0.0
 
 
 def get_mult(
@@ -36,22 +36,23 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
             in_size_local = v_in.index_map.size_local * v_in.block_size
             out_size_local = v_out.index_map.size_local * v_out.block_size
             attach_working_array(mat)  # Ensure working arrays are attached
+            m = typing.cast(typing.Any, mat)
             if transpose:
                 # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                mat._row_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                mat._row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                mat._col_vec.array[:out_size_local] = 0.0
-                mat.mult(mat._row_vec, mat._col_vec, transpose=True)
-                v_out.array[:out_size_local] = mat._col_vec.array[:out_size_local]
+                m._row_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                m._row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                m._col_vec.array[:out_size_local] = 0.0
+                m.mult(m._row_vec, m._col_vec, transpose=True)
+                v_out.array[:out_size_local] = m._col_vec.array[:out_size_local]
             else:
                 # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                mat._row_vec.array[:out_size_local] = 0
-                mat._col_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                mat._col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                mat.mult(mat._col_vec, mat._row_vec)
-                v_out.array[:out_size_local] = mat._row_vec.array[:out_size_local]
+                m._row_vec.array[:out_size_local] = 0
+                m._col_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                m._col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                m.mult(m._col_vec, m._row_vec)
+                v_out.array[:out_size_local] = m._row_vec.array[:out_size_local]
             v_out.scatter_forward()
 
         return mult
@@ -74,25 +75,20 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
         raise TypeError("Matrix type not supported. Expected dolfinx.la.MatrixCSR or PETSc.Mat, got {type(mat)=}.")
 
 
-def _get_interpolation_matrix(
+def _build_interpolation_matrix(
     space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
 ) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
-    """Retrieve or compute the interpolation matrix for a pair of spaces."""
-    key = (id(space_from), id(space_to))
-
-    if key not in _INTERPOLATION_MATRIX_CACHE:
-        if use_petsc:
-            mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to)
-            mat.assemble()
-        else:
-            mat = dolfinx.fem.interpolation_matrix(space_from, space_to)
-            mat.scatter_reverse()
-            # The built in interpolation matrix requires two working arrays
-            attach_working_array(mat)
-
-        _INTERPOLATION_MATRIX_CACHE[key] = mat
+    """Assemble the interpolation matrix for a pair of spaces."""
+    if use_petsc:
+        petsc_mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to)
+        petsc_mat.assemble()
+        return petsc_mat
 
-    return _INTERPOLATION_MATRIX_CACHE[key]
+    mat = dolfinx.fem.interpolation_matrix(space_from, space_to)
+    mat.scatter_reverse()
+    # The built in interpolation matrix requires two working arrays
+    attach_working_array(mat)
+    return mat
 
 
 class InterpolationBlock(Block):
@@ -116,15 +112,27 @@ def __init__(
         self._adj_output: dolfinx.fem.Function | None = None
         self._tlm_output: dolfinx.fem.Function | None = None
         self._hessian_output: dolfinx.fem.Function | None = None
-        self._recompute_output: dolfinx.fem.Function | None = None
+
+        # Interpolation matrix is fixed for the lifetime of this block (tied to
+        # self.space_from/self.space_to), so cache it per-instance rather than
+        # in a module-level dict keyed by id() (which can collide once a
+        # FunctionSpace is garbage collected).
+        self._interpolation_matrix: dolfinx.la.MatrixCSR | "PETSc.Mat" | None = None
 
     def __str__(self):
         return "interpolate_function"
 
+    def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
+        if self._interpolation_matrix is None:
+            self._interpolation_matrix = _build_interpolation_matrix(
+                self.space_from, self.space_to, use_petsc=self._use_petsc
+            )
+        return self._interpolation_matrix
+
     # --- Adjoint ---
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
         adj_input = adj_inputs[0]
@@ -143,7 +151,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
     # --- Tangent Linear Model (TLM) ---
 
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None):
         tlm_input = tlm_inputs[0]
@@ -165,7 +173,7 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
     # --- Hessian ---
 
     def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies):
-        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
+        return self._get_interpolation_matrix()
 
     def evaluate_hessian_component(
         self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None
@@ -191,12 +199,17 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
     def recompute_component(self, inputs, block_variable, idx, prepared):
         func_from = inputs[0]
 
-        if self._recompute_output is None:
-            self._recompute_output = dolfinx.fem.Function(self.space_to)
-
-        self._recompute_output.interpolate(func_from)
-        self._recompute_output.x.scatter_forward()
-
-        # Overload the object to ensure PyAdjoint tracks it properly
-        output = create_overloaded_object(self._recompute_output)
+        # Update the tape's actual output object in-place (rather than a
+        # separately cached Function) so that any Python reference the user
+        # is holding to the interpolated Function stays in sync after the
+        # tape is replayed, matching FunctionAssignBlock's convention. Note
+        # this only holds until the output is first used as a dependency of
+        # another block: pyadjoint then freezes a private checkpoint copy
+        # (see OverloadedType._ad_will_add_as_dependency), so the live
+        # Python object can still go stale after that point. This is an
+        # existing pyadjoint/dolfinx_adjoint-wide characteristic (identical
+        # behavior in FunctionAssignBlock), not specific to interpolation.
+        output = block_variable.saved_output
+        output.interpolate(func_from)
+        output.x.scatter_forward()
         return output

From 1de33910556ece18e6335973c8665e122918e0d4 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Fri, 21 Aug 2026 13:45:15 +0000
Subject: [PATCH 2/5] Restore cross-block interpolation matrix cache, made safe
 against id() reuse

The previous commit removed the shared matrix cache entirely (making it
per-instance on InterpolationBlock) to close the id()-collision hole,
but that gives up the redundant-assembly/MPI-communicator-exhaustion
protection the cache was originally for, since every new block (e.g.
each optimization iteration that creates fresh spaces) reassembles its
own matrix from scratch.

Bring back a shared, module-level cache, but keyed safely this time:
each FunctionSpace that contributes to a cache key gets a
weakref.finalize callback that purges every entry mentioning its id.
Finalizers run at the point the object is actually deallocated, which
is necessarily before CPython can hand that id out again, so a cache
hit always corresponds to spaces that are still alive - no stale
matrix can be served, and entries don't accumulate forever for
short-lived spaces.

Add regression tests: same-space-pair reuse, cache entries getting
purged once their space is garbage collected, and no unbounded growth
across many transient spaces. Verified these tests fail against a
naive id()-keyed cache without the weakref purge (i.e. they would have
caught the original bug).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 src/dolfinx_adjoint/blocks/interpolation.py | 52 ++++++++++++----
 tests/test_interpolate.py                   | 67 ++++++++++++++++++++-
 2 files changed, 107 insertions(+), 12 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index e26a5f2..1c0cf2d 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing
+import weakref
 from typing import Callable
 
 import dolfinx
@@ -10,6 +11,33 @@
 if typing.TYPE_CHECKING:
     from petsc4py import PETSc
 
+# Cache of assembled interpolation matrices, keyed by (id(space_from), id(space_to),
+# use_petsc), shared across all InterpolationBlocks to avoid redundant matrix assembly
+# (and, for PETSc matrices, MPI communicator exhaustion) when the same pair of spaces is
+# interpolated between many times (e.g. across optimization iterations).
+#
+# Keying on id() would normally risk a stale/wrong entry once a FunctionSpace is garbage
+# collected and its id is reused by an unrelated object (ids are only unique among
+# simultaneously-alive objects). We close that hole with weakref.finalize: each space
+# that contributes to a cache key gets a finalizer that purges every entry mentioning
+# its id. Finalizers run at the point the object is actually deallocated, which is
+# necessarily before CPython can hand that id out again, so a cache hit always
+# corresponds to spaces that are still alive.
+_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "dolfinx.la.MatrixCSR | PETSc.Mat"] = {}
+_CACHE_KEYS_BY_SPACE_ID: dict[int, set[tuple[int, int, bool]]] = {}
+
+
+def _purge_cache_entries_for(space_id: int) -> None:
+    for key in _CACHE_KEYS_BY_SPACE_ID.pop(space_id, ()):
+        _INTERPOLATION_MATRIX_CACHE.pop(key, None)
+
+
+def _register_cache_key(space: dolfinx.fem.FunctionSpace, key: tuple[int, int, bool]) -> None:
+    space_id = id(space)
+    if space_id not in _CACHE_KEYS_BY_SPACE_ID:
+        weakref.finalize(space, _purge_cache_entries_for, space_id)
+    _CACHE_KEYS_BY_SPACE_ID.setdefault(space_id, set()).add(key)
+
 
 def attach_working_array(mat: dolfinx.la.MatrixCSR):
     """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication."""
@@ -91,6 +119,18 @@ def _build_interpolation_matrix(
     return mat
 
 
+def _get_interpolation_matrix(
+    space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
+) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
+    """Retrieve or assemble the interpolation matrix for a pair of spaces, cached for reuse."""
+    key = (id(space_from), id(space_to), use_petsc)
+    if key not in _INTERPOLATION_MATRIX_CACHE:
+        _INTERPOLATION_MATRIX_CACHE[key] = _build_interpolation_matrix(space_from, space_to, use_petsc=use_petsc)
+        _register_cache_key(space_from, key)
+        _register_cache_key(space_to, key)
+    return _INTERPOLATION_MATRIX_CACHE[key]
+
+
 class InterpolationBlock(Block):
     """Block for interpolating a dolfinx.fem.Function into another space."""
 
@@ -113,21 +153,11 @@ def __init__(
         self._tlm_output: dolfinx.fem.Function | None = None
         self._hessian_output: dolfinx.fem.Function | None = None
 
-        # Interpolation matrix is fixed for the lifetime of this block (tied to
-        # self.space_from/self.space_to), so cache it per-instance rather than
-        # in a module-level dict keyed by id() (which can collide once a
-        # FunctionSpace is garbage collected).
-        self._interpolation_matrix: dolfinx.la.MatrixCSR | "PETSc.Mat" | None = None
-
     def __str__(self):
         return "interpolate_function"
 
     def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
-        if self._interpolation_matrix is None:
-            self._interpolation_matrix = _build_interpolation_matrix(
-                self.space_from, self.space_to, use_petsc=self._use_petsc
-            )
-        return self._interpolation_matrix
+        return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
 
     # --- Adjoint ---
 
diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py
index 1328412..1106f8c 100644
--- a/tests/test_interpolate.py
+++ b/tests/test_interpolate.py
@@ -1,3 +1,5 @@
+import gc
+
 from mpi4py import MPI
 
 import dolfinx
@@ -7,7 +9,12 @@
 import ufl
 
 from dolfinx_adjoint import Function, assemble_scalar, interpolate
-from dolfinx_adjoint.blocks.interpolation import InterpolationBlock
+from dolfinx_adjoint.blocks.interpolation import (
+    _CACHE_KEYS_BY_SPACE_ID,
+    _INTERPOLATION_MATRIX_CACHE,
+    InterpolationBlock,
+    _get_interpolation_matrix,
+)
 
 # Dynamically determine available matrix backends
 petsc_options = [False]
@@ -148,3 +155,61 @@ def u_ex(mod, x_coords):
     assert np.isclose(min_rate, 3.0, rtol=1e-3, atol=1e-3)
 
     pyadjoint.get_working_tape().clear_tape()
+
+
+# ==============================================================================
+# Test 3: Interpolation Matrix Cache
+# ==============================================================================
+
+
+def test_interpolation_matrix_cache_reused_for_same_spaces(mesh_1D):
+    """The same (space_from, space_to) pair should reuse one assembled matrix."""
+    V_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+    V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    mat1 = _get_interpolation_matrix(V_from, V_to)
+    mat2 = _get_interpolation_matrix(V_from, V_to)
+    assert mat1 is mat2
+
+    # A different space pair must not share the cached matrix.
+    V_other = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 3))
+    mat3 = _get_interpolation_matrix(V_from, V_other)
+    assert mat3 is not mat1
+
+
+def test_interpolation_matrix_cache_purged_when_space_is_garbage_collected(mesh_1D):
+    """Regression test for the id()-collision bug: a cache entry keyed on
+    id(space) must be dropped once that space is garbage collected, otherwise
+    a later, unrelated FunctionSpace object that happens to reuse the same id
+    could silently be served a stale matrix built for a completely different
+    pair of spaces.
+    """
+    V_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+
+    def build_and_return_id():
+        # V_from only lives inside this function, so it is eligible for
+        # collection as soon as it returns.
+        v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+        _get_interpolation_matrix(v_from, V_to)
+        return id(v_from)
+
+    space_id = build_and_return_id()
+    gc.collect()
+
+    assert space_id not in _CACHE_KEYS_BY_SPACE_ID
+    assert all(key[0] != space_id for key in _INTERPOLATION_MATRIX_CACHE)
+
+
+def test_interpolation_matrix_cache_does_not_leak_across_transient_spaces(mesh_1D):
+    """The cache must not grow without bound as short-lived FunctionSpaces
+    (e.g. created once per optimization iteration) are used and discarded.
+    """
+    baseline = len(_INTERPOLATION_MATRIX_CACHE)
+
+    for _ in range(20):
+        v_from = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1))
+        v_to = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 2))
+        _get_interpolation_matrix(v_from, v_to)
+
+    gc.collect()
+    assert len(_INTERPOLATION_MATRIX_CACHE) <= baseline + 1

From f2120e8a2e054a4a5ba54bc94775fd66932bb15b Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Fri, 21 Aug 2026 13:51:42 +0000
Subject: [PATCH 3/5] Replace monkey-patched matrix attributes with a typed
 workspace class

attach_working_array() dynamically attached _row_vec/_col_vec onto
dolfinx.la.MatrixCSR instances, which mypy can't see, so every access
needed typing.cast(Any, ...). Replace it with _MatrixCSRWorkspace, a
small class built once alongside the matrix that holds the working
vectors as real, statically-typed fields, threaded through get_mult(),
_build_interpolation_matrix(), and the cache in place of the bare
matrix. No behavior change (the full working arrays are still zeroed
before each multiply); no more casts or dynamic attributes.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 src/dolfinx_adjoint/blocks/interpolation.py | 73 +++++++++++----------
 1 file changed, 37 insertions(+), 36 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index 1c0cf2d..dcd0c23 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -23,7 +23,7 @@
 # its id. Finalizers run at the point the object is actually deallocated, which is
 # necessarily before CPython can hand that id out again, so a cache hit always
 # corresponds to spaces that are still alive.
-_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "dolfinx.la.MatrixCSR | PETSc.Mat"] = {}
+_INTERPOLATION_MATRIX_CACHE: dict[tuple[int, int, bool], "_MatrixCSRWorkspace | PETSc.Mat"] = {}
 _CACHE_KEYS_BY_SPACE_ID: dict[int, set[tuple[int, int, bool]]] = {}
 
 
@@ -39,48 +39,51 @@ def _register_cache_key(space: dolfinx.fem.FunctionSpace, key: tuple[int, int, b
     _CACHE_KEYS_BY_SPACE_ID.setdefault(space_id, set()).add(key)
 
 
-def attach_working_array(mat: dolfinx.la.MatrixCSR):
-    """Attach working arrays to a dolfinx.la.MatrixCSR for efficient matrix-vector multiplication."""
-    # mat._row_vec/_col_vec are monkey-patched on; cast to Any so mypy doesn't
-    # flag the dynamic attributes.
-    m = typing.cast(typing.Any, mat)
-    if not hasattr(m, "_row_vec"):
-        m._row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
-    if not hasattr(m, "_col_vec"):
-        m._col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
-    m._row_vec.array[:] = 0.0
-    m._col_vec.array[:] = 0.0
+class _MatrixCSRWorkspace:
+    """A dolfinx.la.MatrixCSR paired with pre-allocated working vectors.
+
+    dolfinx.la.MatrixCSR.mult requires vectors built from its own index maps
+    as scratch space. Wrapping them here (built once, alongside the matrix)
+    means they can be reused across repeated matrix-vector multiplications
+    without allocating on every adjoint/TLM/Hessian evaluation, and without
+    monkey-patching dynamic attributes onto the matrix object itself.
+    """
+
+    def __init__(self, mat: dolfinx.la.MatrixCSR):
+        self.mat = mat
+        self.row_vec = dolfinx.la.vector(mat.index_map(0), mat.block_size[0], dtype=mat.data.dtype)
+        self.col_vec = dolfinx.la.vector(mat.index_map(1), mat.block_size[1], dtype=mat.data.dtype)
 
 
 def get_mult(
-    mat: "PETSc.Mat" | dolfinx.la.MatrixCSR,
-    transpose: bool = False,  # type: ignore
+    mat: "PETSc.Mat" | _MatrixCSRWorkspace,
+    transpose: bool = False,
 ) -> Callable[[dolfinx.la.Vector, dolfinx.la.Vector], None]:
     """Return a function that performs matrix-vector multiplication with the given matrix."""
-    if isinstance(mat, dolfinx.la.MatrixCSR):
+    if isinstance(mat, _MatrixCSRWorkspace):
+        workspace = mat
 
         def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
-            # Need to use vectors from
             in_size_local = v_in.index_map.size_local * v_in.block_size
             out_size_local = v_out.index_map.size_local * v_out.block_size
-            attach_working_array(mat)  # Ensure working arrays are attached
-            m = typing.cast(typing.Any, mat)
+            # Zero the full working arrays (including ghosts) before each use
+            # to prevent double-counting in parallel.
+            workspace.row_vec.array[:] = 0.0
+            workspace.col_vec.array[:] = 0.0
             if transpose:
-                # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                m._row_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                m._row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                m._col_vec.array[:out_size_local] = 0.0
-                m.mult(m._row_vec, m._col_vec, transpose=True)
-                v_out.array[:out_size_local] = m._col_vec.array[:out_size_local]
+                workspace.row_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                workspace.row_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                workspace.col_vec.array[:out_size_local] = 0.0
+                workspace.mat.mult(workspace.row_vec, workspace.col_vec, transpose=True)
+                v_out.array[:out_size_local] = workspace.col_vec.array[:out_size_local]
             else:
-                # Prevent double-counting in parallel by zeroing ghosts of input vector
                 # Calculate the exact number of local degrees of freedom
-                m._row_vec.array[:out_size_local] = 0
-                m._col_vec.array[:in_size_local] = v_in.array[:in_size_local]
-                m._col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
-                m.mult(m._col_vec, m._row_vec)
-                v_out.array[:out_size_local] = m._row_vec.array[:out_size_local]
+                workspace.row_vec.array[:out_size_local] = 0
+                workspace.col_vec.array[:in_size_local] = v_in.array[:in_size_local]
+                workspace.col_vec.scatter_forward()  # Ensure ghost values are updated before multiplication
+                workspace.mat.mult(workspace.col_vec, workspace.row_vec)
+                v_out.array[:out_size_local] = workspace.row_vec.array[:out_size_local]
             v_out.scatter_forward()
 
         return mult
@@ -105,7 +108,7 @@ def mult(v_in: dolfinx.la.Vector, v_out: dolfinx.la.Vector):
 
 def _build_interpolation_matrix(
     space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
-) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
+) -> _MatrixCSRWorkspace | "PETSc.Mat":
     """Assemble the interpolation matrix for a pair of spaces."""
     if use_petsc:
         petsc_mat = dolfinx.fem.petsc.interpolation_matrix(space_from, space_to)
@@ -114,14 +117,12 @@ def _build_interpolation_matrix(
 
     mat = dolfinx.fem.interpolation_matrix(space_from, space_to)
     mat.scatter_reverse()
-    # The built in interpolation matrix requires two working arrays
-    attach_working_array(mat)
-    return mat
+    return _MatrixCSRWorkspace(mat)
 
 
 def _get_interpolation_matrix(
     space_from: dolfinx.fem.FunctionSpace, space_to: dolfinx.fem.FunctionSpace, use_petsc: bool = False
-) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
+) -> _MatrixCSRWorkspace | "PETSc.Mat":
     """Retrieve or assemble the interpolation matrix for a pair of spaces, cached for reuse."""
     key = (id(space_from), id(space_to), use_petsc)
     if key not in _INTERPOLATION_MATRIX_CACHE:
@@ -156,7 +157,7 @@ def __init__(
     def __str__(self):
         return "interpolate_function"
 
-    def _get_interpolation_matrix(self) -> dolfinx.la.MatrixCSR | "PETSc.Mat":
+    def _get_interpolation_matrix(self) -> _MatrixCSRWorkspace | "PETSc.Mat":
         return _get_interpolation_matrix(self.space_from, self.space_to, use_petsc=self._use_petsc)
 
     # --- Adjoint ---

From 34a78d1b41441ad7c825a45ed81ac7da4742b699 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Fri, 21 Aug 2026 16:33:48 +0000
Subject: [PATCH 4/5] Fix follow-up review findings from the func_into_space
 merge

- Add the missing Jh_c(c) reset before derivative()/hessian() in the
  Constant Taylor test, mirroring the fix already applied to the
  Function variant (currently masked by the separate Constant
  downcast bug, but confirmed to fail on its own once that's fixed).
- Fix import order in the public interpolate() wrapper.
- Make ExprInterpolationBlock.evaluate_tlm_component re-derive the
  dependency index via get_dependency_index instead of assuming
  tlm_inputs is positionally aligned with self._deps, consistent with
  every other evaluate_*_component method in the class.
- Make ExprInterpolationBlock.recompute_component update
  block_variable.saved_output in place, matching
  InterpolationBlock.recompute_component's convention instead of a
  separately cached buffer wrapped in a fresh overloaded object.
---
 src/dolfinx_adjoint/blocks/interpolation.py | 23 ++++++++++++---------
 src/dolfinx_adjoint/interpolation.py        |  5 +++--
 tests/test_interpolate.py                   |  1 +
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index 4e631ce..551d4ee 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -9,7 +9,6 @@
 import scifem
 import ufl
 from pyadjoint import Block, OverloadedType
-from pyadjoint.overloaded_type import create_overloaded_object
 from pyadjoint.tape import stop_annotating
 from ufl.algorithms.analysis import traverse_unique_terminals
 
@@ -312,7 +311,6 @@ def __init__(
         self._adj_output: dict[int, dolfinx.fem.Function] = {}
         self._tlm_output: dolfinx.fem.Function | None = None
         self._hessian_output: dict[int, dolfinx.fem.Function] = {}
-        self._recompute_output: dolfinx.fem.Function | None = None
 
     def __str__(self):
         return f"interpolate_expression_{str(self.expr)}_to_{str(self.space_to)}"
@@ -380,11 +378,16 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         # Reset output vector to prepare for accumulation
         out_func.x.array[:] = 0.0
 
-        # The TLM is the sum of the Jacobians applied to each perturbation
-        for dep_idx, tlm_input in enumerate(tlm_inputs):
+        # The TLM is the sum of the Jacobians applied to each perturbation. tlm_inputs is
+        # aligned with self.get_dependencies(), not necessarily with self._deps, so the
+        # dependency index into `prepared` is re-derived explicitly rather than assumed
+        # positional, matching every other evaluate_*_component method in this class.
+        for i, dep_bv in enumerate(self.get_dependencies()):
+            tlm_input = tlm_inputs[i]
             if tlm_input is None:
                 continue
 
+            dep_idx = get_dependency_index(self._deps, dep_bv)
             mat = prepared[dep_idx]
             mult = get_mult(mat, transpose=False, accumulate=True)
             mult(tlm_input.x, out_func.x)
@@ -486,15 +489,15 @@ def prepare_recompute_component(self, inputs, relevant_outputs):
         return None
 
     def recompute_component(self, inputs, block_variable, idx, prepared):
-        if self._recompute_output is None:
-            self._recompute_output = dolfinx.fem.Function(self.space_to)
-
         replace_map = {self._deps[i]: inputs[i] for i in range(len(self._deps))}
         updated_expr = ufl.replace(self.expr, replace_map)
 
+        # Update the tape's actual output object in-place, matching
+        # InterpolationBlock.recompute_component and FunctionAssignBlock's convention.
+        output = block_variable.saved_output
         with stop_annotating():
             compiled_expr = dolfinx.fem.Expression(updated_expr, get_interpolation_points(self.space_to))
-            self._recompute_output.interpolate(compiled_expr)
-            self._recompute_output.x.scatter_forward()
+            output.interpolate(compiled_expr)
+            output.x.scatter_forward()
 
-        return create_overloaded_object(self._recompute_output)
+        return output
diff --git a/src/dolfinx_adjoint/interpolation.py b/src/dolfinx_adjoint/interpolation.py
index a6ccb80..b7d61ac 100644
--- a/src/dolfinx_adjoint/interpolation.py
+++ b/src/dolfinx_adjoint/interpolation.py
@@ -1,8 +1,9 @@
 import dolfinx
+import ufl
 from pyadjoint.overloaded_type import create_overloaded_object
 from pyadjoint.tape import annotate_tape, get_working_tape, stop_annotating
-import ufl
-from .blocks.interpolation import InterpolationBlock, ExprInterpolationBlock
+
+from .blocks.interpolation import ExprInterpolationBlock, InterpolationBlock
 from .compat import get_interpolation_points
 
 
diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py
index bc4902e..948b258 100644
--- a/tests/test_interpolate.py
+++ b/tests/test_interpolate.py
@@ -328,6 +328,7 @@ def test_expr_interpolation_taylor_test_constant(mesh_2D, use_petsc):
     min_rate = pyadjoint.taylor_test(Jh_c, c, dc)
     assert np.isclose(min_rate, 2.0, rtol=1e-2, atol=1e-2)
 
+    Jh_c(c)
     dJdm_c = Jh_c.derivative()._ad_dot(dc)
     hessian_c = Jh_c.hessian(dc)
     dHddu_c = hessian_c._ad_dot(dc)

From f4e011e436e20bb9f0f2f5936c55b67ca4ebb0ac Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Fri, 21 Aug 2026 16:49:48 +0000
Subject: [PATCH 5/5] Fix Constant downcast, lazy scifem import, and
 dependency-index reliance

- Function._ad_mul/_ad_add/_ad_copy/_ad_create_checkpoint all hardcoded
  construction of a plain overloaded Function, silently downcasting a
  Constant control during Taylor tests/optimization (surfaced as a
  Control-type TypeError). Add _ad_new_like(), which bypasses the
  constructor via __new__ + Function.__init__ to preserve the concrete
  subclass while still sharing the original function space.
- Make scifem import lazy in blocks/interpolation.py: only
  ExprInterpolationBlock's Jacobian/Hessian assembly needs it, so
  importing dolfinx_adjoint or using the linear InterpolationBlock no
  longer requires it to be installed. Add it as an optional dependency
  group in pyproject.toml (pulled in by the test extra).
- ExprInterpolationBlock's evaluate_*_component/prepare_evaluate_*
  methods no longer re-derive dependency indices via
  get_dependency_index: pyadjoint already hands back the correct index
  through idx and relevant_dependencies' (idx, block_variable) tuples,
  and self._deps is populated in lockstep with get_dependencies() in
  __init__, so position i means the same dependency in both. The
  helper itself is kept for external identity lookups (e.g. tests
  locating a specific control), and the test that manually drives the
  block is fixed to pass relevant_dependencies in pyadjoint's real
  (idx, block_variable) tuple shape instead of a bare block_variable
  list.
---
 pyproject.toml                              |  5 +-
 src/dolfinx_adjoint/blocks/interpolation.py | 57 ++++++++++++---------
 src/dolfinx_adjoint/types/function.py       | 27 ++++++++--
 tests/test_interpolate.py                   |  4 +-
 4 files changed, 61 insertions(+), 32 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 0c7abe4..cf59b1d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,8 @@ dependencies = [
 
 
 [project.optional-dependencies]
-test = ["pytest"]
+scifem = ["scifem"]
+test = ["pytest", "dolfinx-adjoint[scifem]"]
 dev = ["pdbpp", "ipython", "mypy", "ruff"]
 docs = [
       "jupyter-book<2.0",
@@ -29,7 +30,7 @@ docs = [
       "networkx",
       "pygraphviz"
 ]
-all = ["dolfinx-adjoint[test]", "dolfinx-adjoint[dev]", "dolfinx-adjoint[docs]"]
+all = ["dolfinx-adjoint[test]", "dolfinx-adjoint[dev]", "dolfinx-adjoint[docs]", "dolfinx-adjoint[scifem]"]
 
 [tool.pytest.ini_options]
 addopts = ["--import-mode=importlib"]
diff --git a/src/dolfinx_adjoint/blocks/interpolation.py b/src/dolfinx_adjoint/blocks/interpolation.py
index 551d4ee..6a62daf 100644
--- a/src/dolfinx_adjoint/blocks/interpolation.py
+++ b/src/dolfinx_adjoint/blocks/interpolation.py
@@ -6,7 +6,6 @@
 
 import dolfinx
 import dolfinx.fem.petsc
-import scifem
 import ufl
 from pyadjoint import Block, OverloadedType
 from pyadjoint.tape import stop_annotating
@@ -17,6 +16,18 @@
 if typing.TYPE_CHECKING:
     from petsc4py import PETSc
 
+
+def _import_scifem():
+    """Import scifem lazily: only ExprInterpolationBlock needs it, so importing
+    dolfinx_adjoint (or using InterpolationBlock) must not require it to be installed.
+    """
+    try:
+        import scifem
+    except ImportError as e:
+        raise ImportError("scifem is required to interpolate a UFL expression: pip install scifem") from e
+    return scifem
+
+
 # Cache of assembled interpolation matrices, keyed by (id(space_from), id(space_to),
 # use_petsc), shared across all InterpolationBlocks to avoid redundant matrix assembly
 # (and, for PETSc matrices, MPI communicator exhaustion) when the same pair of spaces is
@@ -302,6 +313,11 @@ def __init__(
         self.space_to = func_to.function_space
         self._use_petsc = petsc_mat
 
+        # self._deps and self.get_dependencies() are always populated in lockstep by this
+        # loop, so position i means the same dependency in both: get_dependency_index is
+        # only needed for identity lookups from outside this alignment (e.g. tests locating
+        # a specific control), not for indices pyadjoint already hands back via idx or
+        # relevant_dependencies in the evaluate_*/prepare_evaluate_* methods below.
         self._deps = []
         for op in traverse_unique_terminals(self.expr):
             if isinstance(op, OverloadedType):
@@ -330,6 +346,7 @@ def _assemble_jacobian(self, idx: int, inputs: list | None = None):
         du = ufl.TrialFunction(V_in)
         dE = ufl.derivative(current_expr, target_dep, du)
         dE = ufl.algorithms.apply_derivatives.apply_derivatives(dE)
+        scifem = _import_scifem()
         if self._use_petsc:
             mat = scifem.petsc_interpolation_matrix(dE, self.space_to)
         else:
@@ -341,19 +358,17 @@ def _assemble_jacobian(self, idx: int, inputs: list | None = None):
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
         matrices = {}
-        for dep in relevant_dependencies:
-            idx = get_dependency_index(self._deps, dep)
+        for idx, _dep in relevant_dependencies:
             matrices[idx] = self._assemble_jacobian(idx, inputs)
         return matrices
 
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None):
         adj_input = adj_inputs[0]
-        dep_idx = get_dependency_index(self._deps, block_variable) if block_variable is not None else idx
-        mat = prepared[dep_idx]
+        mat = prepared[idx]
 
-        if dep_idx not in self._adj_output:
-            self._adj_output[dep_idx] = dolfinx.fem.Function(self._deps[dep_idx].function_space)
-        out_func = self._adj_output[dep_idx]
+        if idx not in self._adj_output:
+            self._adj_output[idx] = dolfinx.fem.Function(self._deps[idx].function_space)
+        out_func = self._adj_output[idx]
 
         out_func.x.array[:] = 0.0
         mult = get_mult(mat, transpose=True, accumulate=False)
@@ -365,8 +380,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar
 
     def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs):
         matrices = {}
-        for dep in self.get_dependencies():
-            idx = get_dependency_index(self._deps, dep)
+        for idx in range(len(self._deps)):
             matrices[idx] = self._assemble_jacobian(idx, inputs)
         return matrices
 
@@ -379,15 +393,12 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar
         out_func.x.array[:] = 0.0
 
         # The TLM is the sum of the Jacobians applied to each perturbation. tlm_inputs is
-        # aligned with self.get_dependencies(), not necessarily with self._deps, so the
-        # dependency index into `prepared` is re-derived explicitly rather than assumed
-        # positional, matching every other evaluate_*_component method in this class.
-        for i, dep_bv in enumerate(self.get_dependencies()):
-            tlm_input = tlm_inputs[i]
+        # aligned with self.get_dependencies(), which is aligned 1:1 with self._deps (see
+        # the comment in __init__), so its position doubles as the index into `prepared`.
+        for dep_idx, tlm_input in enumerate(tlm_inputs):
             if tlm_input is None:
                 continue
 
-            dep_idx = get_dependency_index(self._deps, dep_bv)
             mat = prepared[dep_idx]
             mult = get_mult(mat, transpose=False, accumulate=True)
             mult(tlm_input.x, out_func.x)
@@ -424,9 +435,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
                 dE_total = term if dE_total is None else dE_total + term
 
         # 3. Assemble both the Jacobian and the Hessian matrix for each dependency
-        for dep in relevant_dependencies:
-            idx = get_dependency_index(self._deps, dep)
-
+        for idx, _dep in relevant_dependencies:
             # Matrix 1: The standard Jacobian (J)
             J_mat = self._assemble_jacobian(idx, inputs)
 
@@ -447,6 +456,7 @@ def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_
                     if not isinstance(d2E, (int, float)):
                         args = ufl.algorithms.extract_arguments(d2E)
                         if len(args) > 0:
+                            scifem = _import_scifem()
                             if self._use_petsc:
                                 H_mat = scifem.petsc_interpolation_matrix(d2E, self.space_to)
                             else:
@@ -462,12 +472,11 @@ def evaluate_hessian_component(
         hessian_input = hessian_inputs[0]
         adj_input = adj_inputs[0]  # The incoming adjoint sensitivity (y*)
 
-        dep_idx = get_dependency_index(self._deps, block_variable) if block_variable is not None else idx
-        J_mat, H_mat = prepared[dep_idx]
+        J_mat, H_mat = prepared[idx]
 
-        if dep_idx not in self._hessian_output:
-            self._hessian_output[dep_idx] = dolfinx.fem.Function(self._deps[dep_idx].function_space)
-        out_func = self._hessian_output[dep_idx]
+        if idx not in self._hessian_output:
+            self._hessian_output[idx] = dolfinx.fem.Function(self._deps[idx].function_space)
+        out_func = self._hessian_output[idx]
 
         # Reset output vector to 0.0 before accumulation
         out_func.x.array[:] = 0.0
diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py
index 165c36c..70a5114 100644
--- a/src/dolfinx_adjoint/types/function.py
+++ b/src/dolfinx_adjoint/types/function.py
@@ -10,7 +10,6 @@
 from pyadjoint.overloaded_type import (
     FloatingType,
     create_overloaded_object,
-    get_overloaded_class,
     register_overloaded_type,
 )
 from pyadjoint.tape import annotate_tape, get_working_tape, no_annotations, stop_annotating
@@ -69,7 +68,12 @@ def _ad_init_object(cls, obj):
 
     @no_annotations
     def _ad_create_checkpoint(self):
-        checkpoint = create_overloaded_object(self.copy())
+        # Note: self.copy() (dolfinx.fem.Function.copy) always returns a plain
+        # dolfinx.fem.Function regardless of self's concrete type, so wrapping it with
+        # create_overloaded_object would silently downcast a Constant checkpoint to a
+        # plain Function. Use _ad_new_like() instead to preserve the concrete subclass.
+        checkpoint = self._ad_new_like()
+        checkpoint.x.array[:] = self.x.array[:]
         checkpoint.name = self.name + "_checkpoint"
         return checkpoint
 
@@ -109,16 +113,29 @@ def _ad_dot(self, other: typing.Self, options: typing.Optional[dict] = None):
         else:
             raise NotImplementedError("Unknown Riesz representation %s" % riesz_representation)
 
+    def _ad_new_like(self) -> typing.Self:
+        """Create a new, zero-valued instance sharing this object's exact overloaded type and
+        function space.
+
+        Constructing via ``type(self)(...)`` directly does not work here because subclasses
+        such as ``Constant`` take a different constructor signature (mesh and value, not a
+        function space). Going through ``__new__`` and ``Function.__init__`` bypasses that
+        constructor while still producing an instance of the correct concrete subclass.
+        """
+        r = type(self).__new__(type(self), self.function_space)  # type: ignore[call-arg]
+        Function.__init__(r, self.function_space)
+        return r
+
     @no_annotations
     def _ad_mul(self, other: typing.Union[int, float]) -> typing.Self:
         """Multiplication of self with integer or floating value."""
-        r = get_overloaded_class(dolfinx.fem.Function)(self.function_space)
+        r = self._ad_new_like()
         r.x.array[:] = self.x.array * other
         return r
 
     @no_annotations
     def _ad_add(self, other: typing.Self) -> typing.Self:
-        r = get_overloaded_class(dolfinx.fem.Function)(self.function_space)
+        r = self._ad_new_like()
         r.x.array[:] = self.x.array[:] + other.x.array[:]
         return r
 
@@ -188,7 +205,7 @@ def _ad_to_list(m):
 
     def _ad_copy(self):
         """Create a (deep) copy of the function."""
-        r = get_overloaded_class(dolfinx.fem.Function)(self.function_space)
+        r = self._ad_new_like()
         assign(self, r)
         return r
 
diff --git a/tests/test_interpolate.py b/tests/test_interpolate.py
index 948b258..640f8e7 100644
--- a/tests/test_interpolate.py
+++ b/tests/test_interpolate.py
@@ -203,8 +203,10 @@ def test_expr_interpolation_adjoint_property(mesh_2D, use_petsc):
     aligned_tlm_inputs[u_idx] = u
 
     # --- Test Tangent Linear Model and Adjoint ---
+    # relevant_dependencies mirrors pyadjoint's real contract (Block.evaluate_adj): a list of
+    # (idx, block_variable) tuples, where idx is the position in block.get_dependencies().
     mat_tlm = block.prepare_evaluate_tlm(aligned_inputs, aligned_tlm_inputs, None)
-    mat_adj = block.prepare_evaluate_adj(aligned_inputs, [v], [u_bv])
+    mat_adj = block.prepare_evaluate_adj(aligned_inputs, [v], [(u_idx, u_bv)])
 
     tlm_output = block.evaluate_tlm_component(
         inputs=aligned_inputs, tlm_inputs=aligned_tlm_inputs, block_variable=None, idx=0, prepared=mat_tlm