Repository navigation
Differentiate the PDE solution interface in reverse mode - #698
Draft
utkuyilmaz1903 wants to merge 1 commit into
Draft
utkuyilmaz1903 wants to merge 1 commit into
utkuyilmaz1903 wants to merge 1 commit into
Conversation
utkuyilmaz1903
force-pushed
the
solution-interface-ad
branch
from
September 17, 2026 13:32
92ff563 to
c3977b6
Compare
utkuyilmaz1903
force-pushed
the
solution-interface-ad
branch
from
September 28, 2026 09:30
c3977b6 to
c40e559
Compare
utkuyilmaz1903
marked this pull request as ready for review
September 28, 2026 15:14
utkuyilmaz1903
force-pushed
the
solution-interface-ad
branch
from
September 28, 2026 21:42
c40e559 to
ff8df71
Compare
Comment on lines
+158
to
+185
| function ChainRulesCore.rrule( | ||
| config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, | ||
| A::SciMLBase.PDETimeSeriesSolution{T, N, S, D}, | ||
| args::Vararg{Union{Num, Number, AbstractArray, Colon}}; dv = nothing | ||
| ) where {T, N, S, D <: MOLMetadata} | ||
| y = A(args...; dv) | ||
| Δargs = map(_ -> NoTangent(), args) | ||
| function call_pullback(Δ) | ||
| Δ = unthunk(Δ) | ||
| Δ isa AbstractZero && return (NoTangent(), Δargs...) | ||
| Δsol = NoTangent() | ||
| if dv === nothing | ||
| grid_args = _grid_args(A, args) | ||
| for (dvi, Δi) in zip(A.dvs, Δ) | ||
| (Δi === nothing || Δi isa AbstractZero) && continue | ||
| Δsol = Δsol + interpolation_cotangent( | ||
| config, A, dvi, _dv_args(A, dvi, grid_args), Δi | ||
| ) | ||
| end | ||
| else | ||
| cdv = indexed_dependent_variable(A, dv) | ||
| cdv === nothing && throw(not_implemented(dv)) | ||
| Δsol = interpolation_cotangent(config, A, cdv, _grid_args(A, args, cdv), Δ) | ||
| end | ||
| return (Tangent{typeof(A)}(; original_sol = Δsol), Δargs...) | ||
| end | ||
| return y, call_pullback | ||
| end |
Member
There was a problem hiding this comment.
This is piracy. Also, it hits more pieces. How does this do with NeuralPDE.jl's solution?
Member
Author
|
Moving the rules up behind a hook as discussed on Slack; will revisit this PR after that lands. |
utkuyilmaz1903
marked this pull request as draft
September 29, 2026 12:25
This was referenced Sep 30, 2026
utkuyilmaz1903
force-pushed
the
solution-interface-ad
branch
4 times, most recently
from
October 18, 2026 15:57
c4c99ac to
b69db1c
Compare
utkuyilmaz1903
force-pushed
the
solution-interface-ad
branch
from
October 20, 2026 19:35
b69db1c to
e757fcf
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements the reverse-mode hooks of SciML/SciMLBase.jl#1630 for MOL, in place of the ChainRules rules this PR first added here. Depends on that PR: the SciMLBase compat is bumped to 3.59, so CI resolves once that release is out.
The hooks map a cotangent of
sol[u(t, x)]orsol(t, x; dv)back to the solver's solution, for time-dependent and stationary solutions. A grid point that is an unknown takes its share from the state vector by index, and the observed ones, such as eliminated boundary points, go through one observed function of all of them, the way #718 builds the wrapper. On a 21×21 compiled grid that makes a Zygote gradient 4x faster than going through the observed function of the whole field, and Mooncake no longer overflows the stack there. An interpolated value goes to the corners of its cell, and the call hook also returns the coordinate cotangents.test/SolutionAD(lts, 1) checks fields, slices, repeated indices, interpolation at points and lines, coordinates, a time-last field and parameters in observed boundary values against ForwardDiff:Under Mooncake, parameters set with
setp_oopneed SciML/ModelingToolkit.jl#5288. With it the neural network tutorial trains with Mooncake throughsol[u(t, x)], so the docs move from Zygote to Mooncake once it is released. Enzyme is blocked by EnzymeAD/Enzyme.jl#3855.