test/mtk: use Enzyme.autodiff with Duplicated#1456
Merged
ChrisRackauckas merged 1 commit intoMay 26, 2026
Merged
Conversation
5919a89 to
d3f3ddf
Compare
Refactors the DAE-init/SCCNonlinearProblem `@test_broken` block in `test/mtk.jl` to the same pattern as the desauty SCC init test in `desauty_dae_mwe.jl` (SciML#1454): a plain top-level function (no captured mutable closure) with the `ODEProblem` passed as `Duplicated` and the tunables as `Duplicated(t, dt)`. `sensealg` is now a module-level `const _MTK_SENSEALG` (rather than a captured closure variable). Under the rewrite shape, passing it as a function argument under `set_runtime_activity(Reverse)` causes Enzyme to promote it to `Duplicated`, which the `solve_up` Enzyme rule in `DiffEqBaseEnzymeExt` rejects (the rule signature requires `sensealg::Union{Const{Nothing}, Const{<:AbstractSensitivityAlgorithm}}`). Referencing it via a `const` makes Enzyme observe it as `Const`. Status: the block remains `@test_broken`. The first configuration (`prob_correctu0` + `CheckInit`) advances past the closure-capture shadowing issue but still errors during reverse autodiff with the same `MixedDuplicated`/MTK runtime-activity wrapping pattern tracked in SciMLSensitivity.jl#1359 — the structural ODEProblem fields carrying `MTKParameters` are seen as `Duplicated` and propagate into sensealg/u0/p in the underlying `solve_up` dispatch. The refactor matches the documented user-side pattern so that when SciML#1359 lifts, only flipping `@test_broken` → `@test` is needed. The companion Mooncake `@test_broken` block below is unchanged. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
d3f3ddf to
d12ba61
Compare
3 tasks
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.
Summary
Refactors the DAE-init / SCCNonlinearProblem
@test_brokenblock intest/mtk.jlto the sameDuplicated(prob, diprob)pattern as theSCC init test in
desauty_dae_mwe.jl(#1454): a plain top-levelfunction (no captured-mutable closure) with the
ODEProblempassedexplicitly as
Duplicated, and the tunables passed asDuplicated(tunables, dtunables).sensealgis now a module-levelconst _MTK_SENSEALGrather than aclosure capture. Passing it as a function argument under
set_runtime_activity(Reverse)causes Enzyme to promote it toDuplicated, which thesolve_upEnzyme rule inDiffEqBaseEnzymeExt(NewSciML) rejects (its signature requiressensealg::Union{Const{Nothing}, Const{<:AbstractSensitivityAlgorithm}}).Referencing the sensealg through a module-level
constmakes Enzymeobserve it as
Constat the rule dispatch.Status: still
@test_brokenLocal verification:
(prob_correctu0, CheckInit())(the "source of truth"configuration, first entry in
setups): FD gradient computes[-1.91e6, 1.42e6, -9.47e6]. With the rewrite, Enzyme reverseautodiff still errors with
MethodError: no method matching augmented_primal(...)against
DiffEqBase.solve_up— the dispatch resolvedDuplicated{GaussAdjoint}andDuplicated{ODEProblem{...}}(along with
Duplicated{Vector{Float64}}andDuplicated{MTKParameters}foru0/p), but the rule'ssignature requires
sensealg::Const. This is the same familyof
MixedDuplicated/ runtime-activity wrapping issues forMTK-System / NonlinearSolution types tracked in
SciMLSensitivity.jl#1359 (and EnzymeMutabilityException differentiating NonlinearSolve.solve over a RuntimeGeneratedFunction problem (Julia 1.12) EnzymeAD/Enzyme.jl#3117).
Because the first (DAE) configuration still errors after the
rewrite, none of the 15 setups can flip to
@test; the block stays@test_broken. The other configurations were not benchmarkedindependently — local Enzyme runs on this stack are ~12-15 minutes
of compile per configuration, and the structural error reproduces
on the source-of-truth entry, so further per-config exhaustion
wouldn't change the verdict.
This PR is a pure refactor that aligns the test with the
documented user-side pattern (matching #1454's shape) so that when
#1359 lifts, only flipping
@test_broken→@testis needed.Configurations in the loop (unchanged)
15 setups exercise
prob_correctu0(CheckInit),prob_incorrectu0(BrownFullBasicInit / DefaultInit /
nothing),prob_timedepu0(BrownFullBasicInit / DefaultInit /
nothing),prob_correctu0(BrownFullBasicInit / DefaultInit / NoInit /
nothing), andprob_overdetermined(BrownFullBasicInit / DefaultInit / NoInit /nothing). All remain inside one@test_brokenblock.Test plan
test/mtk.jlunder CI on Julia 1.12 with Enzyme 0.13.150and confirm the rewritten
@test_brokenblock still tripsBroken(notPass, which would prompt flipping it to@test).@test_broken(Mooncake) is unaffected.Notes
This is a draft. Please ignore until reviewed by @ChrisRackauckas.
Stacked on top of #1454 (
enzyme-init-test-user-alias-break, alreadymerged into master) — that PR introduced the same pattern and the
NonlinearSolveBase v2.27 compat bump that this PR relies on.
🤖 Generated with Claude Code