Differentiation through MTK DAE initialization with SCCNonlinearProblem is broken for all AD backends

[ChrisRackauckas-Claude]

Description

Mooncake crashes with a MooncakeRuleCompilationError (caused by stack overflow at the LLVM/JIT level) when attempting to differentiate through solve for an MTK ODEProblem that has DAE initialization (algebraic constraints with guesses).

The stack overflow occurs during Mooncake.build_rrule compilation, not during the actual gradient computation. The complex nested type structure of MTK's closure (containing ODEProblem, MTKParameters, GeneratedFunctionWrapper, OverrideInitData, System, etc.) causes Mooncake's rule compiler to overflow.

Partial fixes attempted in ModelingToolkit.jl PR SciML/ModelingToolkit.jl#4331:

• tangent_type(::Type{<:Base.ImmutableDict}) = NoTangent — fixes one recursion path
• tangent_type(::Type{<:System}) = NoTangent — System is self-referential (has systems::Vector{System}, parent::Union{Nothing, System})
• tangent_type(::Type{<:GeneratedFunctionWrapper}) = NoTangent
• tangent_type(::Type{<:ObservedFunctionCache}) = NoTangent
• @from_rrule bridges for MTKParameters constructor, remake_buffer, SetInitialUnknowns

These fix the Julia-level StackOverflowError in tangent_type, but the LLVM-level stack overflow during rule compilation persists.

MWE

import Pkg
Pkg.activate(; temp = true)
Pkg.add(["ModelingToolkit", "OrdinaryDiffEq", "Mooncake",
    "SciMLSensitivity", "SciMLStructures", "SymbolicIndexingInterface"])

using ModelingToolkit, OrdinaryDiffEq, Mooncake
using ModelingToolkit: t_nounits as t, D_nounits as D
using SciMLSensitivity
import SciMLStructures as SS
using SymbolicIndexingInterface

# Minimal DAE system: ODE + algebraic constraint requiring initialization
@parameters a b
@variables x(t) y(t)

eqs = [
    D(x) ~ a * x + y,
    0 ~ x^2 - b * y,  # algebraic constraint
]

@mtkbuild sys = ODESystem(eqs, t)

prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [a => -0.5, b => 2.0],
    guesses = [y => 1.0])
tunables, repack, _ = SS.canonicalize(SS.Tunable(), parameter_values(prob))

# Forward solve works fine
sol = solve(prob, Rodas5P())

# Mooncake gradient fails
loss = let prob = prob, repack = repack
    p -> begin
        new_prob = remake(prob; p = repack(p))
        sol = solve(new_prob, Rodas5P(); abstol = 1e-8, reltol = 1e-6)
        sum(sol)
    end
end

rule = Mooncake.build_rrule(loss, tunables)  # StackOverflow here
val, (_, grad) = Mooncake.value_and_gradient!!(rule, loss, tunables)

Error

Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
MooncakeRuleCompilationError: an error occurred while Mooncake was compiling a rule to differentiate something.

The crash occurs at the LLVM/JIT level during code generation for the complex closure type.

Environment

• Julia 1.12
• ModelingToolkit v11
• Mooncake v0.4
• SciMLSensitivity latest

Related

• Add Mooncake and EnzymeCore extensions for AD compatibility ModelingToolkit.jl#4331 (partial tangent_type fixes)
• Update Core8 MTK tests for mutation-aware AD #1356 (Core8 test updates with @test_broken)

[ChrisRackauckas-Claude]

Full error output

Full output (including precompilation, stack traces): https://gist.github.com/ChrisRackauckas-Claude/8c3c44ee12a5062344d920eccae554d2

Key error excerpt

Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
FAILED: Mooncake.MooncakeRuleCompilationError
MooncakeRuleCompilationError: an error occurred while Mooncake was compiling a rule to differentiate something.

Stacktrace

Stacktrace:
 [1] build_rrule(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx, Mooncake.ReverseMode}, sig_or_mi::Type; debug_mode::Bool, silence_debug_messages::Bool)
   @ Mooncake ~/.julia/packages/Mooncake/hDR3G/src/interpreter/reverse_mode.jl:1164
 [2] build_rrule(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx, Mooncake.ReverseMode}, sig_or_mi::Type)
   @ Mooncake ~/.julia/packages/Mooncake/hDR3G/src/interpreter/reverse_mode.jl:1105
 [3] build_rrule(::Function, ::Vararg{Any}; kwargs::@Kwargs{})
   @ Mooncake ~/.julia/packages/Mooncake/hDR3G/src/interpreter/reverse_mode.jl:1071
 [4] build_rrule(::Function, ::Vector{Float64})
   @ Mooncake ~/.julia/packages/Mooncake/hDR3G/src/interpreter/reverse_mode.jl:1069
 [5] top-level scope
   @ /tmp/mwe_mooncake_full.jl:39

The error includes a very long type signature for Mooncake.build_rrule(Mooncake.MooncakeInterpreter(Mooncake.ReverseMode), Tuple{...}) showing the deeply nested MTK types (GeneratedFunctionWrapper, ObservedFunctionCache, OverrideInitData, InitializationMetadata, etc.) that are captured in the closure. See the gist for the full type signature.

The stack overflow occurs at the LLVM/JIT level during code generation, not at the Julia level. The process sometimes crashes entirely rather than throwing a catchable exception.

[ChrisRackauckas-Claude]

Pure Mooncake MWE (no SciML dependencies)

The root cause is that Mooncake's tangent_type recurses infinitely on any self-referential type. This is a Mooncake bug, not specific to SciML. MTK's System type is self-referential (has systems::Vector{System}, parent::Union{Nothing, System}), which triggers this.

Minimal reproduction with just Mooncake:

using Mooncake

# Any self-referential type triggers this
mutable struct LList{T}
    value::T
    next::Union{Nothing, LList{T}}
end

# tangent_type recurses infinitely
Mooncake.tangent_type(LList{Float64})  # => StackOverflowError

# build_rrule also fails because it needs tangent_type for captured types
ll = LList(1.0, LList(2.0, nothing))
f = let ll = ll
    x -> x[1] * ll.value + x[2] * ll.next.value
end
rule = Mooncake.build_rrule(f, [3.0, 4.0])  # => MooncakeRuleCompilationError (stack overflow)

# Same issue with Base.ImmutableDict (used internally by MTK)
Mooncake.tangent_type(Base.ImmutableDict{Symbol, Any})  # => StackOverflowError

Output:

tangent_type(LList{Float64}):
Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
  FAILED: StackOverflowError

build_rrule on closure capturing LList:
Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
  FAILED: Mooncake.MooncakeRuleCompilationError

tangent_type(ImmutableDict{Symbol,Any}):
Warning: detected a stack overflow; program state may be corrupted, so further execution might be unreliable.
  FAILED: StackOverflowError

This should be filed upstream at https://github.com/chalk-lab/Mooncake.jl. The workaround from MTK PR #4331 is Mooncake.tangent_type(::Type{<:System}) = NoTangent, but this needs to be done for every self-referential type encountered.

[ChrisRackauckas-Claude]

Update: Progress with tangent_type overrides in MTKMooncakeExt

We've confirmed that adding tangent_type overrides for self-referential/structural types in the MTKMooncakeExt extension fixes all tangent_type StackOverflow issues:

# In MTKMooncakeExt (from PR SciML/ModelingToolkit.jl#4331):
Mooncake.tangent_type(::Type{<:Base.ImmutableDict}) = NoTangent
Mooncake.tangent_type(::Type{<:System}) = NoTangent
Mooncake.tangent_type(::Type{<:GeneratedFunctionWrapper}) = NoTangent
Mooncake.tangent_type(::Type{<:ObservedFunctionCache}) = NoTangent

tangent_type results — ALL PASS:

• f.f (GeneratedFunctionWrapper): OK => NoTangent
• f.mass_matrix (UniformScaling): OK => NoTangent
• f.observed (ObservedFunctionCache): OK => NoTangent (was FAIL/StackOverflow without ext)
• f.sys (System): OK => NoTangent (was FAIL/StackOverflow without ext)
• f.initialization_data (OverrideInitData): OK => Tangent{...} (was FAIL/StackOverflow without ext)
• MTKParameters: OK
• ODEFunction: OK
• ODEProblem: OK

build_rrule: NEW error (no longer StackOverflow)

However, build_rrule still fails with a different error during rule compilation:

MethodError: no method matching tangent_type(
  ::Type{Union{Mooncake.NoFData, Mooncake.FData{@NamedTuple{var"##field#754"}}, Mooncake.FData{@NamedTuple{var"##field#755"}}}},
  ::Type{Union{Mooncake.NoRData, Mooncake.RData{@NamedTuple{var"##field#754"}}, Mooncake.RData{@NamedTuple{var"##field#755"}}}}
)

This is a Mooncake-internal issue: tangent_type is being called with 2 arguments (FData and RData types) instead of 1. This appears to be triggered by a Union{...} field type in the OverrideInitData metadata (specifically missing_guess_value which uses a OneOf union type). This is a different bug from the self-referential types issue (chalk-lab/Mooncake.jl#602).

Summary

1. The MTKMooncakeExt tangent_type overrides are correct and necessary
2. There is an additional upstream Mooncake bug with Union-typed fields during FData/RData construction
3. Both issues need to be resolved in Mooncake before MTK DAE problems can be differentiated with Mooncake

[ChrisRackauckas]

Looks like it's an upstream issue chalk-lab/Mooncake.jl#1000

[ChrisRackauckas-Claude]

Status update: Tested on latest versions

Tested the MTK DAE differentiation on Mooncake v0.5.8, ModelingToolkit v11.11.1, SciMLSensitivity v7.96.0, Julia 1.11.9.

Mooncake #1000 (Union-typed fields) — Fixed

The Moshi @data Union-typed field issue (chalk-lab/Mooncake.jl#1000) is resolved in Mooncake v0.5. The MWE passes. Issue closed.

MTK DAE differentiation — Working with non-Mooncake backends

The practical question is whether we can differentiate through MTK DAE solve. Here are the results:

Approach	Result
ForwardDiffSensitivity()	WORKS — gradient [3.46, -0.865]
Zygote default adjoint	WORKS — gradient [1.30, -0.324]
InterpolatingAdjoint(autojacvec=ReverseDiffVJP())	WORKS — gradient [1.30, -0.324]
InterpolatingAdjoint(autojacvec=MooncakeVJP())	FAILS — MooncakeRuleCompilationError
Direct Mooncake.build_rrule(loss, p) on full loss	HANGS — 2+ hours at 100% CPU, no crash

(ForwardDiff result differs because it computes sensitivities differently via forward-mode through the solver internals vs adjoint methods.)

Mooncake-specific issues remain

With Mooncake v0.5, the original StackOverflow is gone — no crash. However:

1. MooncakeVJP adjoint fails: Mooncake can't compile rules for f(u, p, t) when the ODEFunction captures complex MTK types (System, OverrideInitData, InitializationMetadata, etc.). Gets MooncakeRuleCompilationError.

2. Direct build_rrule hangs: Attempting to differentiate the entire loss function through the solver leads to infinite/extremely slow compilation (2+ hours, 100% CPU, ~3GB RAM stable — not a crash, just won't finish).

These are Mooncake compilation scalability issues with deeply nested/complex type trees, not SciMLSensitivity bugs per se.

Conclusion

The SciMLSensitivity issue with MTK DAE differentiation is effectively resolved for practical use — Zygote (default adjoint), ReverseDiffVJP, and ForwardDiff all work correctly. The remaining Mooncake-specific failures are upstream Mooncake compilation scalability problems with MTK's type complexity.

Recommend closing this issue since the core use case (differentiating through MTK DAE solve) works. The Mooncake-specific path can be tracked upstream if needed.

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

[ChrisRackauckas-Claude]

Correction: SCCNonlinearProblem differentiation is broken

My previous comment was inaccurate — I tested a simple DAE that produced a regular NonlinearProblem for initialization, not an SCCNonlinearProblem. When testing with the De Sauty bridge circuit from test/desauty_dae_mwe.jl (which actually produces an SCCNonlinearProblem for initialization), all gradient methods fail.

Test setup

Tested on ModelingToolkit v11.11.1, SciMLSensitivity v7.96.0, OrdinaryDiffEq v6.99.2, Julia 1.11.9.

using ModelingToolkitStandardLibrary.Electrical
using ModelingToolkitStandardLibrary.Blocks: Sine

# De Sauty bridge circuit (from SciMLSensitivity test suite)
desauty_model = create_model()
dsys = structural_simplify(desauty_model)
dprob = ODEProblem(dsys, [dsys.resistor1.v => 1.0], (0.0, 0.1))
iprob = dprob.f.initialization_data.initializeprob

typeof(iprob)  # SCCNonlinearProblem — this is the key difference

Results

Differentiating through init problem directly

Method	Result
FiniteDiff	WORKS — nonzero gradients
Zygote + SteadyStateAdjoint	FAILS — SciMLBase.AdjointNotFoundError

The AdjointNotFoundError occurs because concrete_solve.jl explicitly excludes SCCNonlinearProblem:

# From SciMLSensitivity/src/concrete_solve.jl
const ConcreteNonlinearProblem = Union{
    NonlinearProblem, SciMLBase.ImmutableNonlinearProblem, SciMLBase.SteadyStateProblem,
}

The SCCNonlinearSolveChainRulesCoreExt tries to call _concrete_solve_adjoint but falls through to the SciMLBase fallback which throws AdjointNotFoundError.

Differentiating through full ODE solve

Method	Result
FiniteDiff	WORKS — [-30094.56, 0.0, 0.0, 86103.56, 0.0]...
Zygote default	FAILS — MethodError: no method matching getindex(::Nothing, ::Int64)
Zygote + ReverseDiffVJP	FAILS — same getindex(::Nothing, ::Int64) error

The getindex(::Nothing, ::Int64) error occurs in MTK generated code at SymbolicUtils/src/code.jl:1036 during initialization.

Root cause

There are two separate issues:

1. SciMLSensitivity has no adjoint rule for SCCNonlinearProblem — The ConcreteNonlinearProblem union explicitly excludes it, so _concrete_solve_adjoint never dispatches for SCC problems. The comment in the code says SCCNonlinearProblem "should differentiate through their individual NonlinearProblem sub-solves" but this doesn't work in practice.

2. Full ODE solve gradient crashes in MTK generated initialization code — Even without directly differentiating the init problem, the ODE adjoint path hits getindex(::Nothing, ::Int64) in SymbolicUtils-generated code during the initialization phase.

Note

SciMLSensitivity's own test/parameter_initialization.jl has all initialization gradient tests marked @test_broken with the comment: "Known issue: MTK parameter initialization gradients may return zeros due to missing AD overloads."

Conclusion

This issue should remain open. The original Mooncake StackOverflow is fixed in v0.5, but the underlying problem — differentiating through MTK DAE initialization when SCCNonlinearProblem is used (the default) — is broken for all AD backends, not just Mooncake. The title could be updated to reflect the broader scope.

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

[ChrisRackauckas-Claude]

PR #1375 replaces the De Sauty component model test with a flat DAE that tests both use_scc=false (NonlinearProblem) and use_scc=true (SCCNonlinearProblem) initialization paths. All AD gradient tests are currently @test_broken for both modes.