Clone appctx for the cached adjoint replay solvers#5170
Conversation
A bit of a doubling up on the already-present adj_sol on solver blocks.
* reduce nlvs adjoint test size * nlvs adjoint caching - bundle up recompute objects
When a NonlinearVariationalSolver is created with an appctx dict, the cached solvers used for tape replay deep-copy the form coefficients, but the appctx was passed through unchanged (forward solver) or dropped outright (adjoint and TLM solvers). Preconditioners that read from appctx (e.g. MassInvPC for Schur complement Stokes) would see stale end-of-forward-run values instead of the tape-replayed values. The forward cache now exposes its coefficient replace map, and a new _ad_clone_kwargs helper applies ufl.replace to every UFL expression in appctx using that map. Since the tangent and adjoint solvers build their forms from the cloned forward problem, all three cached solvers share the single forward replace map, so their appctx entries point at the same cloned coefficients that update_dependencies keeps in sync with the tape. The tangent and adjoint solvers fall back to the forward kwargs' appctx when their own kwargs do not carry one: their operators are linearisations of the forward operator, so a preconditioner that needs appctx for the forward solve needs the same for them. The appctx pop in solve_init_params is kept: the legacy GenericSolveBlock adjoint path solves with an assembled matrix through firedrake.solve, which rejects an appctx kwarg, so appctx must still be dropped there.
|
Why are there two PRs for this? If this goes into |
|
The two are not the same patch landing twice, so closing this one would drop the fix rather than defer it. #5171 against This PR is the same fix expressed against the world that @JHopeCollins's |
The previous comment claimed the assembled adjoint solve does not accept appctx, which is misleading now that LinearSolver is a LinearVariationalSolver subclass and forwards appctx. The actual rejection is one layer up: the solve(A, x, b) dispatch validates kwargs in _extract_linear_solver_args and refuses a top-level appctx, reading it only from solver_parameters.
sghelichkhani
force-pushed
the
sghelichkhani/appctx-on-cached-solver-block
branch
from
aea71cb to
e9af276
Compare
angus-g
force-pushed
the
JHopeCollins/nlvs-hessian-fix
branch
from
6f6fd15 to
440f386
Compare
4 participants
When an annotated
NonlinearVariationalSolveris created with anappctx, the replay machinery deep-copies the form coefficients but passes the appctx through untouched (forward solver) or drops it entirely (adjoint and TLM solvers). Any preconditioner that reads UFL expressions out of the appctx, MassInvPC being the obvious one for Schur complement Stokes, ends up looking at the user's original coefficients rather than the clones the tape keeps in sync, so during replay it builds the preconditioner from stale end-of-forward-run values.This applies the forward coefficient replace map to UFL entries in the appctx when the cached solvers are built. Since the tangent and adjoint solvers derive their forms from the cloned forward problem, all three share the one map and so point at the same cloned coefficients that
update_dependenciesupdates at replay time. The adjoint and TLM solvers fall back to the forward appctx when their own kwargs don't carry one; David confirmed this is always safe since the adjoint equation uses J^T (https://firedrakeproject.slack.com/archives/C1Q0Y6H8A/p1776171359074139). Theappctxpop insolve_init_paramsstays, because the legacyGenericSolveBlockadjoint path solves with an assembled matrix and can't accept the kwarg; there's a regression test pinning that path too.The new tests in
tests/firedrake/adjoint/test_appctx.pycheck object identity between the appctx and the cached forms, and use a recording preconditioner to verify the values seen during forward and adjoint replay match the forward trajectory.Note this targets the #4638 refactor branch, so it only reaches
mainwhen that merges. The same bug exists in the current release through the old_ad_problem_clonemachinery, so we should probably do a separate PR basing the release branch to fix released Firedrake as well; I have that version of the fix ready on a separate branch.Disclosure: parts of this change were prepared with the help of Claude Code (Opus 4.8).