Skip to content

Commit 7896045

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #1453 from ChrisRackauckas-Claude/cc/enzyme-test-broken-annotations
Annotate @test_broken Enzyme blocks with Const(loss) + set_runtime_activity
2 parents 0499b87 + 3b1e9b3 commit 7896045

3 files changed

Lines changed: 74 additions & 35 deletions

File tree

‎test/desauty_dae_mwe.jl‎

Lines changed: 43 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ using ForwardDiff
1010
using Tracker
1111
using Enzyme
1212
using Mooncake
13+
using NonlinearSolve: NewtonRaphson
1314

1415
# DAE with nonlinear algebraic constraints forming an SCC chain.
1516
# Inspired by the De Sauty bridge DAE but written as a flat system
@@ -120,28 +121,49 @@ eqs = [
120121
end
121122

122123
@testset "Enzyme through init" begin
123-
# Status (verified 2026-04-10 with Enzyme 0.13, NonlinearSolve
124-
# 4.17, SciMLBase 2.153, ModelingToolkit current release):
124+
# Annotations follow the documented user-side pattern:
125+
# `Const(loss)` for the closure that captures the mutable
126+
# `NonlinearProblem`/`SCCNonlinearProblem`, and
127+
# `set_runtime_activity(Reverse)` so Enzyme's activity analysis
128+
# tolerates the runtime-activity transitions through MTK's
129+
# `remake` path. The inner `solve` pins `NewtonRaphson()`
130+
# explicitly so Enzyme's type analysis does not trip on the
131+
# polyalgorithm Union NonlinearSolve would otherwise dispatch
132+
# through. The previously-reported `EnzymeMutabilityException`
133+
# on the mutable closure capture is correct upstream behavior
134+
# per EnzymeAD/Enzyme.jl#3117 — annotating with `Const` is the
135+
# fix.
125136
#
126-
# * Julia 1.10 (LTS): hits an `EnzymeMutabilityException` because
127-
# the closure capturing `iprob`/`irepack` cannot be proven
128-
# read-only. Wrapping `init_loss` in `Const(...)` advances past
129-
# the activity check but then crashes the LLVM GC invariant
130-
# verifier with `Illegal inttoptr` during `MTK.remake`/
131-
# `SciMLStructures.replace`. Even calling
132-
# `Enzyme.set_runtime_activity(Reverse)` produces the same
133-
# LLVM crash. The issue reproduces equally for `use_scc=false`
134-
# and `use_scc=true` and is independent of SCCNonlinearSolve.
135-
# * Julia 1.11+: same crashes plus a separate
136-
# `IllegalTypeAnalysisException` on `Base._typed_vcat!` inside
137-
# SCCNonlinearSolve's solution assembly.
138-
#
139-
# Tracking issues: NonlinearSolve.jl#869, Enzyme.jl#2699,
140-
# Enzyme.jl#3021 (vcat type analysis), and the upstream MTK
141-
# remake/Enzyme interaction.
142-
@test_broken begin
143-
igs = Enzyme.gradient(Enzyme.Reverse, init_loss, itunables)
144-
!iszero(sum(igs))
137+
# With these annotations, the plain `NonlinearProblem` case
138+
# (use_scc = false) now passes. The `SCCNonlinearProblem` case
139+
# (use_scc = true) still trips a `MixedDuplicated` /
140+
# `Core.SimpleVector` MethodError further down in Enzyme's
141+
# runtime-activity wrapping for the MTK-System /
142+
# NonlinearSolution types involved in SCC sub-problem
143+
# assembly — tracked in SciMLSensitivity.jl#1359. When that
144+
# lifts, flipping `@test_broken` → `@test` in the `use_scc`
145+
# branch is the only change needed here.
146+
enzyme_init_loss = let iprob = iprob, irepack = irepack
147+
p -> begin
148+
iprob2 = remake(iprob, p = irepack(p))
149+
sol = solve(iprob2, NewtonRaphson())
150+
sum(sol.u)
151+
end
152+
end
153+
if use_scc
154+
@test_broken begin
155+
igs = Enzyme.gradient(
156+
Enzyme.set_runtime_activity(Enzyme.Reverse),
157+
Enzyme.Const(enzyme_init_loss), itunables,
158+
)
159+
!iszero(sum(igs))
160+
end
161+
else
162+
igs = Enzyme.gradient(
163+
Enzyme.set_runtime_activity(Enzyme.Reverse),
164+
Enzyme.Const(enzyme_init_loss), itunables,
165+
)
166+
@test !iszero(sum(igs))
145167
end
146168
end
147169

‎test/mtk.jl‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -160,11 +160,19 @@ setups = [
160160
]
161161

162162
# Reverse-mode AD through DAE initialization with SCCNonlinearProblem mutation.
163-
# Marked as broken until Enzyme/Mooncake fully support this pattern.
164-
# Enzyme blockers (see NonlinearSolve.jl#869, issue #1358):
165-
# - Julia 1.12: LLVM crash (Enzyme rules disabled by VERSION < v"1.12" guard)
166-
# - Julia 1.10: EnzymeMutabilityException in remake, MixedReturnException with
167-
# default PolyAlgorithm, NamedTuple broadcast error with MTKParameters
163+
# Annotations follow the documented user-side pattern: `Const(loss)` for the
164+
# closure that captures the mutable `ODEProblem`, and
165+
# `set_runtime_activity(Reverse)` so Enzyme's activity analysis tolerates the
166+
# runtime-activity transitions through MTK's `remake` path. The inner solve
167+
# already pins `Rodas5P()` (no polyalgorithm Union for Enzyme's type analysis
168+
# to trip on). The previously-reported `EnzymeMutabilityException` on the
169+
# mutable closure capture is correct upstream behavior per
170+
# EnzymeAD/Enzyme.jl#3117 — annotating with `Const` is the fix. With these
171+
# annotations the chain advances through the activity layer; the remaining
172+
# blocker is a `MixedDuplicated` / `Core.SimpleVector` MethodError further
173+
# down in Enzyme's runtime-activity wrapping for MTK-System /
174+
# NonlinearSolution types — tracked in SciMLSensitivity.jl#1359. When that
175+
# lifts, flipping `@test_broken` → `@test` is the only change needed here.
168176
@test_broken begin
169177
grads = map(setups) do setup
170178
prob, tunables, repack, init = setup
@@ -185,7 +193,10 @@ setups = [
185193
sum(new_sol)
186194
end
187195
end
188-
Enzyme.gradient(Enzyme.Reverse, loss, tunables)
196+
Enzyme.gradient(
197+
Enzyme.set_runtime_activity(Enzyme.Reverse),
198+
Enzyme.Const(loss), tunables,
199+
)
189200
end
190201
all(x ≈ grads[1] for x in grads)
191202
end

‎test/parameter_initialization.jl‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,16 @@ tunables, repack, _ = SS.canonicalize(SS.Tunable(), parameter_values(prob))
7070
end
7171

7272
# Exercises the EnzymeOriginator method of `_init_originator_gradient`
73-
# added alongside this testset. Currently @test_broken because the
74-
# outer Enzyme.gradient over `remake(prob; p = repack(tunables))`
75-
# itself fails with `EnzymeRuntimeActivityError` from MTK's `remake`
76-
# path — same upstream issue tracked by NonlinearSolve.jl#869 /
77-
# Enzyme.jl#2699 / SciMLSensitivity.jl#1415. When that clears, this
78-
# should pass without further changes (the dispatch already routes
79-
# the init step through Enzyme natively).
73+
# added alongside this testset. Annotations follow the documented
74+
# user-side pattern: `Const(loss)` for the closure that captures the
75+
# mutable `ODEProblem`, and `set_runtime_activity(Reverse)` so Enzyme's
76+
# activity analysis tolerates the runtime-activity transitions through
77+
# MTK's `remake` path. With these in place the activity layer is
78+
# handled; the remaining blocker is a `MixedDuplicated` /
79+
# `Core.SimpleVector` MethodError further down in Enzyme's
80+
# runtime-activity wrapping for MTK-System / NonlinearSolution
81+
# types — tracked in SciMLSensitivity.jl#1359. When that lifts,
82+
# flipping `@test_broken` → `@test` is the only change needed here.
8083
@testset "Adjoint through Prob (Enzyme)" begin
8184
sensealg = SciMLSensitivity.GaussAdjoint(
8285
autojacvec = SciMLSensitivity.EnzymeVJP(),
@@ -89,7 +92,10 @@ tunables, repack, _ = SS.canonicalize(SS.Tunable(), parameter_values(prob))
8992
end
9093
end
9194
@test_broken begin
92-
g = Enzyme.gradient(Enzyme.Reverse, Enzyme.Const(loss), copy(tunables))[1]
95+
g = Enzyme.gradient(
96+
Enzyme.set_runtime_activity(Enzyme.Reverse),
97+
Enzyme.Const(loss), copy(tunables),
98+
)[1]
9399
any(!iszero, g)
94100
end
95101
end

0 commit comments

Comments
 (0)