-
Notifications
You must be signed in to change notification settings - Fork 199
Adjoint appctx and projection taping issues
Notes for the meeting on 2026-06-16. Two separate things to raise. The first is a
contained bug fix that mostly needs a decision on strategy. The second is a tangle around
project where the immediate fix is easy but the real question is whether the whole
annotation path wants redesigning, and that overlaps with Josh's NLVS refactor.
When an annotated NonlinearVariationalSolver carries an appctx, the replay machinery
deep-copies the form coefficients so it can drive each tape step with the right values, but
it hands the appctx dictionary through still pointing at the user's original
coefficients. So any preconditioner that reads UFL out of appctx sees stale,
end-of-forward-run values during replay rather than the per-step values the tape keeps in
sync.
The case that drove this out is MassInvPC used as the Schur-complement approximation in a
fieldsplit Stokes solve, with appctx={"mu": viscosity}. In a mantle convection inversion
the viscosity follows an Arrhenius law with a contrast of about 10^5, so the viscosity at an
early timestep and at the end of the run differ by orders of magnitude. The forward pressure
solve converges in 6 to 30 iterations; on replay (any Jhat(m), Jhat.derivative(), or a
Taylor test) the same solve hits the 200 iteration cap and diverges, because the
preconditioner is frozen at the end-of-run viscosity. Beyond the convergence failure it can
quietly degrade the gradient, since the Stokes fields are no longer solved to tolerance at
each step.
The mechanism is just object identity. The replayed forward F and J are built with
deepcopy plus ufl.replace, giving fresh Function objects the tape updates per step. The
appctx dict was never run through that replace map, so appctx["mu"] still references the
original Function, which sits at its end-of-run value. The fix points appctx["mu"] at the
same cloned Function the form uses, and then it is kept in sync for free. Worth knowing: a
Taylor test cannot detect this, because a converged Krylov solve does not depend on its
preconditioner, so the gradient comes out right either way. The real detector is a recorder
PC that logs the value it actually sees at each initialize/update.
The fix exists in two forms because the code it patches exists in two forms.
#5171 targets release. It
repairs the legacy cloning path (_ad_problem_clone and _ad_adj_lvs_problem in
variational_solver.py), threading the coefficient replace maps back out so the appctx can
be cloned alongside the forms. This is the genuine bug fix for released Firedrake and is
ready (mergeable, tests pass, reverting the source makes the identity and replay tests fail
as expected).
#5170 is the same fix expressed
against Josh's nlvs-hessian-fix branch, the head of
#4638. That refactor retires the
two methods #5171 patches and replaces them with the _ad_forward_cache /
_ad_tangent_cache / _ad_adjoint_cache machinery. The two are not the same patch landing
twice. If #5171 reaches main and then #4638 rebases on top, the release edits land on code
that has already been rewritten, so the appctx fix evaporates into dead code and has to be
reinstated against the cache. That is what #5170 does, and it additionally covers the tangent
(TLM) path, which has no counterpart in the legacy structure. Connor asked on the PR why
there are two; this is the answer, and it is the same trap we hit early on, where git happily
auto-merges the old fix into a method the refactor no longer calls.
A nice property of the refactor is that the forward, tangent and adjoint solvers all derive their forms from one cloned forward problem, so there is a single replace map rather than the two independent ones the legacy design has. That makes the appctx clone cleaner there than on release.
Removing the appctx pop from solve_init_params looks like the obvious simplification but
it crashes the legacy solve path. That function is shared with GenericSolveBlock, the live
block for every annotated firedrake.solve(F == 0, u, ...). Its adjoint solve goes through
solve(A, x, b), which dispatches to _la_solve and _extract_linear_solver_args, and that
validates kwargs against a fixed list with no appctx in it, so a top-level appctx raises
Illegal keyword argument 'appctx'. The Navier-Stokes demo uses exactly this pattern. So the
pop stays, and the cached adjoint and tangent solvers instead fall back to the forward
appctx via a default_appctx argument. Carrying appctx into the adjoint solver is always
safe because the adjoint uses J^T; David Ham confirmed this on Slack on 2026-04-14.
The second trap is in the test, not the source. The original test solved twice (mu=1 then mu=10) but only assembled the functional after the second solve, so the first solve's output only fed the second solve's initial guess and carried no adjoint. pyadjoint therefore never ran the first block's adjoint solve and the replay test passed vacuously. Fixed by accumulating the functional after each solve.
On #5171 Connor's suggestion is to make the linear solver path itself accept an appctx
argument, so _la_solve / _extract_linear_solver_args would no longer reject it. Then the
GenericSolveBlock adjoint solve would not choke, the pop would not be needed, and there
would be no default_appctx special-casing on our side. His framing was that the PR makes
sense but he does not want to bake in unneeded technical debt.
That is a cleaner end state, but it is a broader change that touches the linear solver API, which is a bigger thing to put on a release branch than a contained adjoint-only fix. The question I want to settle is the sequencing:
- For
release, do we take the minimal contained fix as it stands in #5171 (lowest risk on a release branch), or do we hold it for the linear-solver-accepts-appctx change? - My inclination is minimal fix into release now, and treat "teach the linear solver about
appctx" as a separate, non-release cleanup on
mainthat would then let us drop the pop and thedefault_appctxfallback in both the legacy and cached code. - For #5170, agree how it folds into Josh's work: keep it as a PR against
nlvs-hessian-fix, or hand the diff to Josh to absorb into #4638 directly. Either way it should not be closed, because the release fix does not survive the refactor on its own.
There are several distinct problems with how project is taped. Naming them so they do not
get conflated:
Solver parameters never recorded (#5172).
An annotated project recorded no solver configuration at all, so every replay and adjoint
ran with the global defaults (preonly plus LU, i.e. mumps) instead of whatever the forward
projection used. A deliberately unconverged Richardson projection replays with a functional
three orders of magnitude off. This was really four faults at once: annotate_project
forwarded only the adjoint callback kwargs and dropped solver_parameters and
form_compiler_parameters; the Function.project path had its own duplicate wrapper that
forwarded nothing; the Projector's defaults are applied inside ProjectorBase.__init__,
which runs after annotation, so even with forwarding fixed an unparameterised project would
still record nothing; and projecting a Function onto its own space shortcuts to an assignment
in the forward while the tape still recorded a solving block.
Direct Projector use is never taped (#5176).
Projector(...).project() is invisible to the tape entirely. Only the project(...)
free function and Function.project are annotated, so a whole entry point silently produces
no adjoint.
Two forward paths, one block. Same-space projection assigns, cross-space projection solves,
but historically both recorded the same solving ProjectBlock, so the tape did not match the
forward model.
#5175 (against main, mergeable)
addresses #5172. It extracts the parameter defaulting out of ProjectorBase.__init__ into a
module-level resolve_projection_solver_parameters, adds a _project_block factory that
records the effective parameters (forward kwargs carry solver plus form-compiler parameters,
adjoint kwargs carry solver only since the mass matrix is self-adjoint and _la_solve
rejects form-compiler parameters), deletes the duplicated Function.project wrapper in
favour of delegating to the now-correct free function, and records a same-space projection as
a FunctionAssignBlock to match the forward. One deliberate behaviour change to flag: an
unparameterised taped projection now replays with the cg/icc mass solve it actually performed
(rtol 1e-8) rather than an exact direct solve, so anything sensitive at that tolerance could
shift.
#5175 works and closes #5172, but it is a record-the-parameters patch and it cannot honour
everything the Projector knows. It cannot capture the slate inverse or the
constant_jacobian choice, and it does nothing for #5176. The annotation path is genuinely
fragmented: three entry points (projection.project, Function.project,
Projector.project), two forward behaviours (solve and assign), and parameters bolted on
after annotation rather than being part of the recorded object.
The structurally honest fix is a block that holds the Projector itself and reuses its
cached solver on replay, the way SupermeshProjectBlock already holds its Projector. That
would tape all three entry points uniformly, naturally pick up the assign-versus-solve split,
and carry the slate-inverse and constant-jacobian fidelity for free. The catch is that this
reshapes adjoint_utils/blocks/solving.py, which is exactly the file Josh's #4638 rewrites,
so a redesign wants to be built on top of that refactor rather than racing it. There is
already a coordination note on #4638 pointing at this.
So the decision I want from the meeting: take #5175 now as the contained bug fix for the solver parameters and keep #5176 open as the redesign, to be built on the CachedSolverBlock machinery once #4638 lands? Or hold for a single Projector-holding block that subsumes both? My inclination is to land #5175 now because the parameter loss is a real correctness bug people are hitting today, and schedule the Projector-holding redesign against #4638.