Annotate @test_broken Enzyme blocks with Const(loss) + set_runtime_activity

The three `@test_broken Enzyme.gradient` blocks (`mtk.jl` line 168,
`parameter_initialization.jl` "Adjoint through Prob (Enzyme)",
`desauty_dae_mwe.jl` "Enzyme through init") previously called
`Enzyme.gradient(Enzyme.Reverse, loss, tunables)` — i.e. relied on
`Enzyme.gradient` to infer annotations on a closure that captures a
mutable `ODEProblem`/`NonlinearProblem`. That triggers
`EnzymeMutabilityException` on a captured-mutable, which per Billy Moses
(EnzymeAD/Enzyme.jl#3117 close) is correct Enzyme behavior, not a bug:
`Enzyme.gradient` can't annotate captures, so the user must do it.

Apply the documented user-side pattern in all three blocks:

  Enzyme.gradient(
      Enzyme.set_runtime_activity(Enzyme.Reverse),
      Enzyme.Const(loss),
      tunables,
  )

Plus, for `desauty_dae_mwe.jl`, pin `solve(iprob2, NewtonRaphson())` so
Enzyme's type analysis does not trip on the polyalgorithm Union the
default NonlinearSolve dispatch would otherwise emit. (`mtk.jl` already
uses `Rodas5P()`; `parameter_initialization.jl` uses an `ODEProblem`
solve with `GaussAdjoint` sensealg, no NonlinearSolve polyalg in the AD
hot path.)

These tests **stay `@test_broken`** — the activity layer is correct
after this change, but the chain still hits a `MixedDuplicated` /
`Core.SimpleVector` MethodError further down in Enzyme's
runtime-activity wrapping for MTK-`System` / `NonlinearSolution` types.
Tracked in #1359, with one upstream piece
already filed at SciML/ModelingToolkit.jl#4553 (declare `AbstractSystem`
as `inactive_type`). When the remaining upstream lifts, flipping
`@test_broken` → `@test` is the only change needed in these blocks.

Refs #1323, #1359; EnzymeAD/Enzyme.jl#3117.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

test/desauty_dae_mwe.jl

@@ -10,6 +10,7 @@ using ForwardDiff
 using Tracker
 using Enzyme
 using Mooncake
+using NonlinearSolve: NewtonRaphson
 
 # DAE with nonlinear algebraic constraints forming an SCC chain.
 # Inspired by the De Sauty bridge DAE but written as a flat system
@@ -120,27 +121,36 @@ eqs = [
         end
 
         @testset "Enzyme through init" begin
-            # Status (verified 2026-04-10 with Enzyme 0.13, NonlinearSolve
-            # 4.17, SciMLBase 2.153, ModelingToolkit current release):
-            #
-            #   * Julia 1.10 (LTS): hits an `EnzymeMutabilityException` because
-            #     the closure capturing `iprob`/`irepack` cannot be proven
-            #     read-only. Wrapping `init_loss` in `Const(...)` advances past
-            #     the activity check but then crashes the LLVM GC invariant
-            #     verifier with `Illegal inttoptr` during `MTK.remake`/
-            #     `SciMLStructures.replace`. Even calling
-            #     `Enzyme.set_runtime_activity(Reverse)` produces the same
-            #     LLVM crash. The issue reproduces equally for `use_scc=false`
-            #     and `use_scc=true` and is independent of SCCNonlinearSolve.
-            #   * Julia 1.11+: same crashes plus a separate
-            #     `IllegalTypeAnalysisException` on `Base._typed_vcat!` inside
-            #     SCCNonlinearSolve's solution assembly.
-            #
-            # Tracking issues: NonlinearSolve.jl#869, Enzyme.jl#2699,
-            # Enzyme.jl#3021 (vcat type analysis), and the upstream MTK
-            # remake/Enzyme interaction.
+            # Annotations follow the documented user-side pattern:
+            # `Const(loss)` for the closure that captures the mutable
+            # `NonlinearProblem`/`SCCNonlinearProblem`, and
+            # `set_runtime_activity(Reverse)` so Enzyme's activity analysis
+            # tolerates the runtime-activity transitions through MTK's
+            # `remake` path. The inner `solve` pins `NewtonRaphson()`
+            # explicitly so Enzyme's type analysis does not trip on the
+            # polyalgorithm Union NonlinearSolve would otherwise dispatch
+            # through. The previously-reported `EnzymeMutabilityException`
+            # on the mutable closure capture is correct upstream behavior
+            # per EnzymeAD/Enzyme.jl#3117 — annotating with `Const` is the
+            # fix. With these annotations the chain advances through the
+            # activity layer; the remaining blocker is a `MixedDuplicated`
+            # / `Core.SimpleVector` MethodError further down in Enzyme's
+            # runtime-activity wrapping for MTK-System / NonlinearSolution
+            # types — tracked in SciMLSensitivity.jl#1359. When that
+            # lifts, flipping `@test_broken` → `@test` is the only change
+            # needed here.
+            enzyme_init_loss = let iprob = iprob, irepack = irepack
+                p -> begin
+                    iprob2 = remake(iprob, p = irepack(p))
+                    sol = solve(iprob2, NewtonRaphson())
+                    sum(sol.u)
+                end
+            end
             @test_broken begin
-                igs = Enzyme.gradient(Enzyme.Reverse, init_loss, itunables)
+                igs = Enzyme.gradient(
+                    Enzyme.set_runtime_activity(Enzyme.Reverse),
+                    Enzyme.Const(enzyme_init_loss), itunables,
+                )
                 !iszero(sum(igs))
             end
         end

test/mtk.jl

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

test/parameter_initialization.jl

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