Support checkpointing of time-dependent adjoints by finsberg · Pull Request #71 · scientificcomputing/dolfinx-adjoint

finsberg · 2026-08-25T10:13:45Z

Taping a time-dependent model keeps every intermediate state alive, because the adjoint sweep needs each of them on the way back, and for a long simulation that is what exhausts memory first. pyadjoint already knows how to trade that memory for recomputation: it consumes schedules from checkpoint_schedules and drives the forward and adjoint sweeps accordingly. What was missing was the DOLFINx side of the contract.

Nothing here reimplements scheduling. Users configure checkpointing entirely through pyadjoint's own API — tape.enable_checkpointing(...) and tape.timestepper(...) — and this PR adds a single name of its own, enable_disk_checkpointing, only because pyadjoint has no equivalent.

The fix

A block must never put a non-overloaded value on the tape. Both solver blocks built their replay solution vector as a plain dolfinx.fem.Function, and that value became the block's output. Outside checkpointing nothing ever asks such a value to checkpoint itself, which is why this went unnoticed; under a schedule a stored output is re-stored on a later pass and it fails. The construction is now extracted into one helper so the rule is stated once rather than duplicated across LinearProblemBlock and NonlinearProblemBlock.

Disk storage

Schedules that store on disk need somewhere to put a function. This uses h5py directly rather than taking on adios4dolfinx, because adios4dolfinx solves a harder problem than we have: its checkpoints are portable across process counts, which requires storing the mesh and dof permutations alongside the values. Checkpointing during an adjoint computation never needs that — the data is written and read within one run, by the same processes, against an unchanged mesh and partition. Under those assumptions the payload is just a process's local values and the format is a flat array.

These are therefore snapshot checkpoints, not a user-facing artefact: they cannot be reopened by a later run or on a different number of processes, and they are deleted automatically. Anyone wanting a checkpoint that outlives the run wants adios4dolfinx.

Two details are load-bearing and both were found by MPI runs rather than by reasoning:

Ghost values are stored alongside the owned ones. Owned values alone would be smaller, but restoring them needs a forward scatter, which is collective. Restoring is not a collective-safe place to be: it happens whenever some block wants a value, filtered by a cache whose lifetime depends on when the garbage collector runs — and the collector does not run at the same moment on every process. This deadlocked on the 54th restore of a two-process run, one rank taking a cached value while the other read and communicated. Storing the ghosts makes restoring purely local and removes the hazard rather than narrowing it.
Checkpoint data stays in one file until teardown rather than rolling over when the tape resets. pyadjoint resets package data before recomputing the forward, but then restores the initial condition from a checkpoint written while taping, so data from before a reset is still live and discarding it there breaks the restore.

use_mpio selects one shared MPI-IO file or one file per process, defaulting to whichever fits the run; it is explicit rather than purely automatic so both layouts are testable on one machine.

Verification

72 passed, 2 xfailed, serially and under mpirun -n 2; mypy and ruff clean.

The tests compare the checkpointed gradient against the un-checkpointed one and run a Taylor test, over both file layouts. The two checks earn their keep separately: gradient equality alone would not catch both paths being equally wrong. Perturbation directions are interpolated analytic expressions rather than random numbers, because they must agree across processes and per-rank random values do not — that alone moved an observed Taylor rate from 1.53 to 2.0.

The demo checks the same properties as it goes, so a regression fails the docs build.

Known limitation, and it is not from this change

NonlinearProblem cannot currently be advanced over timesteps. Its adjoint is wrong — Taylor rate -0.41 where the linear model gives 2.0 — and it fails identically with checkpointing disabled. Confirmed on unmodified main.

The unknown is a coefficient of the residual, so it is registered among the block's own dependencies. Once its incoming value is itself control-dependent — which is exactly what a time loop creates — the adjoint contribution for it is computed against a residual in which that value no longer appears, and ufl.adjoint raises IndexError on the resulting argument-less form. Suppressing that error is not a fix: the gradient is then silently wrong. test_poisson_mother's nonlinear case escapes this because its residual is actually linear in the unknown, so the block appears never to have been exercised with a genuinely nonlinear one.

Recorded as a strict xfail so it reports itself when fixed. Until then NonlinearProblem cannot be covered by the checkpointing tests.

Two upstream fragilities worked around, not fixed

Tape.clear_tape() resets the checkpoint manager but leaves _eagerly_checkpoint_outputs and latest_checkpoint set, so a tape that has once been checkpointed keeps checkpointing outputs eagerly even after being cleared.
LinearProblem writes its PETSc options into the global options database under a fixed default prefix and never removes them, so a later solver constructed without explicit options silently inherits them. This makes test_time_dependent_bc_replay order-dependent: on current main it already fails if run after test_linear_solver, and passes in CI only because of alphabetical ordering.

The new test file isolates both so the suite stays green, but neither is actually fixed.

Not in scope

Reusing solver objects across timesteps, and Hessians under checkpointing. The latter needs an upstream pyadjoint change: Tape.evaluate_adj consults the checkpoint manager but evaluate_tlm and evaluate_hessian do not, so the tangent-linear model runs over a tape whose checkpoints the reverse sweep has already released.

🤖 Generated with Claude Code

pyadjoint already drives checkpoint schedules from checkpoint_schedules; what was missing was the DOLFINx side of the contract. Blocks must not put non-overloaded values on the tape. Both solver blocks built their replay solution vector as a plain dolfinx.fem.Function, which became the block's output. Outside checkpointing nothing asks such a value to checkpoint itself, so this went unnoticed; under a schedule a stored output is re-stored on a later pass and it fails. Extracted the construction so the rule lives in one place. Added a disk backend for schedules that store on disk, written with h5py rather than taking on adios4dolfinx: these are snapshot checkpoints, valid only within the run that wrote them and against an unchanged partition, so the payload is just a process's local values with no mesh or permutation data. Ghost values are stored alongside the owned ones so that restoring needs no communication -- restores are filtered by a cache whose lifetime depends on when the garbage collector runs, which is not the same moment on every process, and a collective call on that path deadlocks. Checkpoint data stays in one file until teardown, because pyadjoint resets package data before recomputing but then restores an initial condition written while taping. enable_disk_checkpointing is the only name this adds; schedules and the timestepping loop stay pure pyadjoint. Tests compare the checkpointed gradient against the un-checkpointed one and run a Taylor test, over both file layouts, serially and on two processes. NonlinearProblem cannot yet be advanced over timesteps: its adjoint is wrong with checkpointing disabled too (Taylor rate -0.41 against 2 for the linear model), which is pre-existing and unrelated. Recorded as a strict xfail so it reports itself when fixed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The disk backend needs h5py, and the CI image does not ship it: both the test module and the demo's disk section failed to import. Declared as a dependency rather than an extra, since the demo exercises it on every docs build. The lazy import stays as a safety net, and an h5py without MPI support still works -- each process then writes its own checkpoint file. Point the demo's API references at the packages they belong to via intersphinx, and cite the checkpointing literature with a per-document key prefix so the labels stay unique across pages. Refer to io4dolfinx rather than its former name. The demo now also turns disk checkpointing off when it is done, which is what deletes the checkpoint files.

jorgensd · 2026-08-25T10:29:44Z

+grad_ckpt = [np.copy(g.x.array) for g in rf_ckpt.derivative()]
+
+assert np.isclose(J_plain, J_ckpt)
+for a, e in zip(grad_ckpt, grad_plain):


Use strict in zip.

Done, in both places in the demo, and in the two equivalent zips in tests/test_checkpointing.py for consistency.

jorgensd · 2026-08-25T10:32:40Z

+directions = []
+for k in range(num_steps):
+    h = dolfinx_adjoint.Function(V, name=f"direction_{k}")
+    # Interpolated rather than random: the direction has to be the same on every process, and
+    # per-process random numbers are not.
+    h.interpolate(lambda x, k=k: np.sin((k + 1) * np.pi * x[0]) * np.cos(np.pi * x[1]))
+    directions.append(h)


Shouldn't annotate=False be set in initialization of these functions, or use with pyadjoint.pause_annotating():.... ?

Right, they were being taped. Wrapped in pyadjoint.stop_annotating() in both the demo and the test helper: the directions are inputs to the test, not part of the model.

Simplify the note on tape.timestepper, and say why Firedrake can pass a bare range: it sets tape.progress_bar, whose iter() returns a real iterator, while the default passes the argument straight through for next() to choke on. Use strict zips and np.testing.assert_allclose so a mismatch reports what differed rather than just failing, and rewrite the sentence introducing the Taylor test to say what it means.

jorgensd · 2026-08-25T10:33:52Z

+
+assert np.isclose(J_plain, J_disk)
+for a, e in zip(grad_disk, grad_plain):
+    assert np.allclose(a, e)


Strict in zip and assert allclose instead of allclose.

Both done in that block already: zip(..., strict=True) and np.testing.assert_allclose.

jorgensd · 2026-08-25T10:35:44Z

+    Overloaded rather than plain, because whatever the block returns from
+    `recompute_component` becomes its output on the tape: under a checkpoint schedule a stored
+    output is asked to checkpoint itself again on a later pass, which a plain
+    `dolfinx.fem.Function` cannot do. Outside checkpointing nothing asks, which is why a plain


Use {py:class}dolfinx.fem.Function here

jorgensd · 2026-08-25T10:36:54Z

+    """
+    with stop_annotating():
+        if isinstance(u, dolfinx.fem.Function):
+            return Function(u.function_space, name=u.name + "_initial_guess")


Function is imported as _Function in line 10.

Good spot — switched the isinstance check to _Function to match the annotation.

jorgensd · 2026-08-25T10:38:13Z


    @no_annotations
    def _ad_create_checkpoint(self):
+        from ..checkpointing import maybe_disk_checkpoint


Can import be moved to the top of the file? Is there any reason not to do it?

Moved to module scope. It was local to avoid a cycle, but checkpointing.py only needs Function for annotations now (under typing.TYPE_CHECKING), so there is no runtime cycle left.

jorgensd · 2026-08-25T10:39:00Z

        return checkpoint

    def _ad_restore_at_checkpoint(self, checkpoint):
+        from ..checkpointing import SnapshotCheckpoint


Same comment regarding import. Why not at top of file?

Same — moved to module scope.

zip(..., strict=True) needs 3.10, so declare the floor rather than leave it implicit. assign_linear_combination reached for .function_space and .x on whatever extract_linear_combination returned, which UFL types as BaseCoefficient. Newer UFL types that tightly enough for mypy to reject it, failing the formatting job on main since before this branch. Assigning a linear combination genuinely needs the DOLFINx Function that carries the degrees of freedom, so check for one and say so, rather than reaching for an attribute UFL does not promise.

jorgensd · 2026-08-25T11:01:42Z

@@ -0,0 +1,139 @@
+# # Time-distributed control
+# Based on example from https://dolfin-adjoint.github.io/dolfin-adjoint/documentation/time-distributed-control/time-distributed-control.html


Why do you need this demo, and the time_distributed_control checkpointing demo. Which should be kept and why?

This remove pushed by accident. I have now removed it

Understood — I picked up your removal when I merged main into the branch, so the file is gone here too. Nothing outstanding on this one.

jorgensd · 2026-08-25T11:02:58Z

+stored alongside the owned ones, which keeps restoring free of communication -- see `_layout`.
+
+Snapshot checkpoints are therefore not portable. They cannot be reopened by a later run, or on a
+different number of processes. For a checkpoint that outlives the run, use ``io4dolfinx``.


Use {py:mod}io4dolfinx.

Done — {py:mod}io4dolfinx`` in the module docstring, and the demo links to the repository.

jorgensd · 2026-08-25T11:04:42Z

+
+
+def _import_h5py():
+    try:


Add note regarding lazy import in python 3.15 https://docs.python.org/3.15/reference/simple_stmts.html#lazy-imports

h5py is a declared dependency now, so the lazy import has gone entirely and it is imported at module scope.

jorgensd · 2026-08-25T11:12:49Z

+    if not shared_file:
+        return n_local, n_local, 0
+    # Collective, but called only from the write path, which every process reaches together.
+    sizes = comm.allgather(n_local)


Replace with exscan.
As the following example shows:

from mpi4py import MPI import time import numpy as np comm = MPI.COMM_WORLD rng = np.random.default_rng(seed=42) z = np.random.randint(0, 100, size=1)[0] start_time = time.perf_counter() vals = comm.allgather(z) offset = np.sum(vals[: comm.rank]) end_time = time.perf_counter() start_ex = time.perf_counter() offset_ex = comm.exscan(z, op=MPI.SUM) if comm.rank == 0: offset_ex = 0 end_ex = time.perf_counter() np.testing.assert_array_equal(offset, offset_ex) print(f"Rank {comm.rank}: offset = {offset}, time taken = {end_time - start_time:.6f} seconds") print(f"Rank {comm.rank}: offset_ex = {offset_ex}, time taken = {end_ex - start_ex:.6f} seconds") print( f"Ratio allreduce/exscan: {end_time - start_time:.6f} / {end_ex - start_ex:.6f} = {(end_time - start_time) / (end_ex - start_ex):.2f}" )

root@docker-desktop:~/shared# mpirun --allow-run-as-root -n 8 python3 mwe_exscan.py Rank 3: offset = 201, time taken = 0.001738 seconds Rank 3: offset_ex = 201, time taken = 0.000080 seconds Ratio allreduce/exscan: 0.001738 / 0.000080 = 21.78 Rank 1: offset = 95, time taken = 0.000637 seconds Rank 1: offset_ex = 95, time taken = 0.000075 seconds Ratio allreduce/exscan: 0.000637 / 0.000075 = 8.47 Rank 6: offset = 330, time taken = 0.001919 seconds Rank 6: offset_ex = 330, time taken = 0.000077 seconds Ratio allreduce/exscan: 0.001919 / 0.000077 = 24.76 Rank 2: offset = 145, time taken = 0.001058 seconds Rank 2: offset_ex = 145, time taken = 0.000079 seconds Ratio allreduce/exscan: 0.001058 / 0.000079 = 13.37 Rank 5: offset = 283, time taken = 0.001528 seconds Rank 5: offset_ex = 283, time taken = 0.000081 seconds Ratio allreduce/exscan: 0.001528 / 0.000081 = 18.83 Rank 4: offset = 254, time taken = 0.001921 seconds Rank 4: offset_ex = 254, time taken = 0.000077 seconds Ratio allreduce/exscan: 0.001921 / 0.000077 = 24.96 Rank 7: offset = 416, time taken = 0.000133 seconds Rank 7: offset_ex = 416, time taken = 0.000075 seconds Ratio allreduce/exscan: 0.000133 / 0.000075 = 1.77 Rank 0: offset = 0.0, time taken = 0.000638 seconds Rank 0: offset_ex = 0, time taken = 0.000073 seconds Ratio allreduce/exscan: 0.000638 / 0.000073 = 8.73

exscan is always faster.

Done — comm.exscan(n_local, op=MPI.SUM) with the rank-0 None handled, plus one allreduce for the total length, which the dataset has to be sized with. Thanks for the benchmark; the prefix-sum-over-allgather was me reaching for the first thing that worked rather than the operation this actually is.

jorgensd · 2026-08-25T11:13:21Z

+    rolling to a new one when the tape resets, and tearing down.
+    """
+
+    def __init__(self, path: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool):


use pathlib.Path here as input instead.

Done — pathlib.Path for the file path, the checkpointer directory, and the dirname argument.

jorgensd · 2026-08-25T11:14:10Z

+        self.path = path
+        self.comm = comm


Store with _path and _comm and make @property decorators to fetch them.
same for shared_file below.

Done — _path, _comm and _shared_file with @property accessors.

jorgensd · 2026-08-25T11:14:56Z

+    rolling to a new one when the tape resets, and tearing down.
+    """
+
+    def __init__(self, path: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool):


Document what input args do. It is unclear what cleanup does.

Done. cleanup now says what it controls: whether the files and the temporary directory are deleted on teardown, with the note that they are unreadable by a later run either way, so keeping them is only useful for debugging.

jorgensd · 2026-08-25T11:17:37Z

+
+    __slots__ = ("_file", "_key", "_space", "_cls", "_n_local", "_offset", "_name", "_cache", "__weakref__")
+
+    def __init__(self, file: _CheckpointFile, key: str, function: dolfinx.fem.Function, n_local: int, offset: int):


Document input arguments.

jorgensd · 2026-08-25T11:18:15Z

+
+    __slots__ = ("_file", "_key", "_space", "_cls", "_n_local", "_offset", "_name", "_cache", "__weakref__")
+
+    def __init__(self, file: _CheckpointFile, key: str, function: dolfinx.fem.Function, n_local: int, offset: int):


Should the input be a dolfinx_adjoint.Function or the parent class dolfinx.Function?

Changed to dolfinx_adjoint.Function — only overloaded types are ever checkpointed, so the parent class was too loose. It is imported under typing.TYPE_CHECKING since dolfinx_adjoint.types imports this module at runtime.

jorgensd · 2026-08-25T11:21:27Z

+class _DiskCheckpointer(TapePackageData):
+    """Tape-attached state owning the checkpoint files for one tape."""
+
+    def __init__(self, directory: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool, owns_directory: bool):


Document inputs.

jorgensd · 2026-08-25T11:24:42Z

+            "use_mpio=True requires an MPI-enabled build of h5py. Use use_mpio=False to write "
+            "one checkpoint file per process instead."
+        )
+    owns_directory = dirname is None


Here we need an all-gather to ensure that all processes owns directory or not, as it can cause a deadlock in _DiskCheckpointer at

def close(self) -> None: """Close the current file and remove the directory if this object created it.""" self._file.close() self._storing = False if self._owns_directory: self._comm.Barrier() if self._comm.rank == 0: try: os.rmdir(self._directory) except OSError: # pragma: no cover - non-empty when cleanup was disabled pass

if self._owns_directory is not synced across all processes.

Good catch, and it is a real deadlock. Fixed: the processes now agree explicitly rather than each deciding for itself.

without_dirname = comm.allreduce(int(dirname is None), op=MPI.SUM) if without_dirname not in (0, comm.size): raise ValueError("dirname must be given on every process or on none of them, ...") owns_directory = without_dirname == comm.size

So a mixed call fails immediately with a clear message instead of hanging in close().

jorgensd · 2026-08-25T11:37:18Z

+        self._storing = False
+
+
+def maybe_disk_checkpoint(function: dolfinx.fem.Function) -> typing.Optional[SnapshotCheckpoint]:


Is this the correct class, did you say that we always want to use dolfinx_adjoint.Function.

Furthermore, use SnapshotCheckpoint | None instead of typing.Optional.

Changed to dolfinx_adjoint.Function, and to SnapshotCheckpoint | None. I have replaced typing.Optional throughout the module while I was there.

jorgensd · 2026-08-25T11:37:42Z

+
+
+def enable_disk_checkpointing(
+    dirname: typing.Optional[str] = None,


Use pathlib.Path.

The parallel test run hung intermittently -- about one run in three locally, and in CI it sat for 42 minutes before being cancelled. Two ranks were in different collectives at once: one inside dolfinx's mpi_jit, the other inside dolfinx.fem.petsc.LinearProblem.__del__, which destroys the KSP and matrices and is collective. Every problem.solve() creates a block owning its own LinearProblem, and pyadjoint blocks sit in reference cycles, so a discarded tape's solvers are freed by the cyclic garbage collector rather than by refcounting. That runs when a process crosses an allocation threshold, which is not the same moment on every process, so one process enters a collective destructor the others are not in. Checkpointing made it likely by recomputing the forward many times over. Collect deliberately at points every process reaches together -- around each test and before building a tape -- so those destructors stay in step. Twelve consecutive parallel runs of the full suite now pass where two in three hung before. The underlying hazard is architectural: a block should not own PETSc objects it cannot destroy deterministically. Sharing one solver across a problem's blocks is the real fix, and is the solver-reuse work already planned separately. Also stop allocating the adjoint right-hand side per evaluation and reuse it, which removes another per-call collective allocation from that path. Review comments: exclusive scan instead of gathering every size, pathlib for paths, private attributes behind properties, documented arguments, cross-reference roles, imports at module scope, dolfinx_adjoint.Function in annotations, and modern optional syntax. Whether the checkpoint directory was created here is now agreed across processes rather than assumed: disagreeing would deadlock teardown, which synchronises before removing it. Perturbation directions are built with annotation stopped; they are inputs to a test, not part of the model.

jorgensd · 2026-08-25T13:02:35Z

+    """
+    with stop_annotating():
+        if isinstance(u, _Function):
+            return Function(u.function_space, name=u.name + "_initial_guess")


Shouldn't this and the next line then use _Function(...) as a constructor.

No — and thank you for asking, because the asymmetry is deliberate and it read like an oversight, which means it needed saying in the code rather than in a thread.

_Function is dolfinx.fem.Function; Function is our overloaded subclass. The check has to be against the base class, because u may be either a plain function or an overloaded one and the overloaded type is a subclass of it. The construction has to be of the overloaded type, because what this returns is what recompute_component hands back, and that becomes the block's output on the tape. A plain dolfinx.fem.Function cannot checkpoint itself, so constructing one here reintroduces exactly the failure this helper exists to prevent — it is the bug this PR opened with.

I have narrowed the return annotation to the overloaded type and spelled the distinction out in the docstring, so the next reader gets the answer without having to ask.

Nonmatching interpolation landed on main while this branch was open, and the resulting conflict stopped GitHub computing a merge ref, which is why the pull-request workflows silently stopped running: only the push-triggered documentation build was left. Both sides added a demo to the table of contents; kept both, with the checkpointing demo next to the one it extends. The index_map property main added is annotated dolfinx.cpp.la.IndexMap, which does not exist -- the type is dolfinx.cpp.common.IndexMap, publicly dolfinx.common.IndexMap. mypy rejects the merged tree without this, so it is corrected here rather than left to fail the formatting job.

The isinstance check is against dolfinx's Function and the construction is of the overloaded one, which reads like an oversight when the two names are an underscore apart. It is not: what arrives is only known to be a dolfinx.fem.Function, because the overloaded type is a subclass and either may be passed, while what leaves ends up on the tape and so must be overloaded. Constructing the plain one would reintroduce the failure this helper exists to prevent. Return annotation narrowed to the overloaded type to match.

finsberg · 2026-08-25T15:01:25Z

+    if _checkpointer is None or not _checkpointer.storing:
+        return None
+    return _checkpointer.store(function)


Don't we need to add a

global _checkpointer

here?

No — global is only needed to rebind a module-level name. maybe_disk_checkpoint only reads _checkpointer, and a read resolves through the module globals at call time, so it always sees the current value. The two functions that assign to it, enable_disk_checkpointing and disable_disk_checkpointing, do both declare global.

Adding it here would be harmless but misleading, since it would suggest this function reassigns the checkpointer when it does not.

Demonstrated against the built module:

before enable : None after enable : _DiskCheckpointer storing=False -> None storing=True -> SnapshotCheckpoint after disable : None -> maybe_disk_checkpoint: None

The disk tests are the standing proof of the same thing: they only pass because maybe_disk_checkpoint picks up the checkpointer that enable_disk_checkpointing installed.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>