Skip to content

test/mtk: use Enzyme.autodiff with Duplicated#1456

Merged
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:enzyme-mtk-test-duplicated
May 26, 2026
Merged

test/mtk: use Enzyme.autodiff with Duplicated#1456
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:enzyme-mtk-test-duplicated

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Summary

Refactors the DAE-init / SCCNonlinearProblem @test_broken block in
test/mtk.jl to the same Duplicated(prob, diprob) pattern as the
SCC init test in desauty_dae_mwe.jl (#1454): a plain top-level
function (no captured-mutable closure) with the ODEProblem passed
explicitly as Duplicated, and the tunables passed as
Duplicated(tunables, dtunables).

sensealg is now a module-level const _MTK_SENSEALG rather than a
closure capture. 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 (NewSciML) rejects (its signature requires
sensealg::Union{Const{Nothing}, Const{<:AbstractSensitivityAlgorithm}}).
Referencing the sensealg through a module-level const makes Enzyme
observe it as Const at the rule dispatch.

Status: still @test_broken

Local verification:

  • Config (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 reverse
    autodiff still errors with
    MethodError: no method matching augmented_primal(...)
    against DiffEqBase.solve_up — the dispatch resolved
    Duplicated{GaussAdjoint} and Duplicated{ODEProblem{...}}
    (along with Duplicated{Vector{Float64}} and
    Duplicated{MTKParameters} for u0/p), but the rule's
    signature requires sensealg::Const. This is the same family
    of MixedDuplicated / runtime-activity wrapping issues for
    MTK-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 benchmarked
independently — 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@test is 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), and
prob_overdetermined (BrownFullBasicInit / DefaultInit / NoInit /
nothing). All remain inside one @test_broken block.

Test plan

  • Run test/mtk.jl under CI on Julia 1.12 with Enzyme 0.13.150
    and confirm the rewritten @test_broken block still trips
    Broken (not Pass, which would prompt flipping it to @test).
  • Confirm the second @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, already
merged into master) — that PR introduced the same pattern and the
NonlinearSolveBase v2.27 compat bump that this PR relies on.

🤖 Generated with Claude Code

@ChrisRackauckas-Claude ChrisRackauckas-Claude force-pushed the enzyme-mtk-test-duplicated branch from 5919a89 to d3f3ddf Compare May 26, 2026 00:35
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>
@ChrisRackauckas-Claude ChrisRackauckas-Claude force-pushed the enzyme-mtk-test-duplicated branch from d3f3ddf to d12ba61 Compare May 26, 2026 01:25
@ChrisRackauckas ChrisRackauckas marked this pull request as ready for review May 26, 2026 02:38
@ChrisRackauckas ChrisRackauckas merged commit abe7a98 into SciML:master May 26, 2026
40 of 50 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants