build_rrule hangs on MTK-generated ODEProblem types (DynamicDerivedRule infinite type inference)

[ChrisRackauckas-Claude]

Summary

Mooncake.build_rrule succeeds for the top-level function, but execution hangs indefinitely during lazy sub-rule compilation (DynamicDerivedRule) when Mooncake encounters ModelingToolkit-generated ODEProblem types. The hang is in abstract_call_gf_by_type → typeinf_edge cycles in Mooncake's abstract interpretation.

This blocks using Mooncake as the outer AD for MTK-based models in SciMLSensitivity.jl (see SciML/SciMLSensitivity.jl#1358).

MWE

# Mooncake v0.5.23, ModelingToolkit v11.17.0, OrdinaryDiffEq v6.108.0, Julia 1.11.9

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using OrdinaryDiffEq
using SciMLStructures: Tunable, replace
using Mooncake

@parameters k
@variables X(t)
@named sys = System([D(X) ~ -k * X], t)
sys = mtkcompile(sys)

prob = ODEProblem(sys, [X => 1.0], (0.0, 1.0), [k => 0.5])

function loss(pval)
    new_p = replace(Tunable(), prob.p, pval)
    sol = solve(prob, Tsit5(); p = new_p, saveat = 0.1)
    return sum(reduce(hcat, sol.u))
end

println("Primal: ", loss([0.5]))  # works

# build_rrule succeeds (~15s):
rule = Mooncake.build_rrule(loss, [0.5])

# But execution hangs here during DynamicDerivedRule lazy compilation:
Mooncake.value_and_pullback!!(rule, 1.0, loss, [0.5])  # hangs

Behavior

• Without MTK (plain ODEProblem{true}(f!, u0, tspan, p)): Works correctly, gradients match ForwardDiff.
• With MTK (even simple D(X) ~ -k*X): build_rrule succeeds in ~15s, but value_and_pullback!! hangs.
• With MTK + DAE initialization (algebraic constraints + guesses): Same hang, but the types are even more complex.

Stack trace (from SIGTERM)

The hang is in deeply nested abstract_call_gf_by_type → typeinf_edge → abstract_call_method cycles:

DynamicDerivedRule → build_rrule → generate_ir → lookup_ir → typeinf_ircode →
  typeinf → typeinf_edge → abstract_call_method → abstract_call_gf_by_type →
    abstract_call_gf_by_type (Mooncake override) → abstract_call_known →
      abstract_call → typeinf_edge → ... (repeats)

This suggests Mooncake's abstract interpretation doesn't terminate for the complex parameterized types used by MTK (RuntimeGeneratedFunctions, GeneratedFunctionWrapper, ObservedFunctionCache, etc.).

Versions

• Mooncake v0.5.23
• ModelingToolkit v11.17.0
• OrdinaryDiffEq v6.108.0
• Julia 1.11.9

[yebai]

The current implementation lacks a differentiation rule for solve.

Mooncake failed to differentiate the following method: kwcall(::NamedTuple, ::typeof(solve), prob::SciMLBase.AbstractDEProblem, args...) @ DiffEqBase ~/.julia/packages/DiffEqBase/XLzeg/src/solve.jl:575

As a result, Mooncake cannot differentiate through calls to solve, which is encountered in DynamicDerivedRule above.

Concretely, the implementation here neither imports nor defines a rule for Core.kwcall, which is required to handle keyword-based calls to solve.

I have noted this in SciML/SciMLSensitivity.jl#1358. Multi-day runs with automated tools (e.g., Claude) are unlikely to resolve it, as addressing the problem requires substantial domain expertise.

[ChrisRackauckas]

The current implementation lacks a differentiation rule for solve.

That is not true, there is a rule precisely for this solve. There has been for months now, but the system isn't hitting it for some reason that keeps getting highlighted.

[ChrisRackauckas]

Concretely, the implementation here neither imports nor defines a rule for Core.kwcall, which is required to handle keyword-based calls to solve.

Do I need to double up all rules for keyword argument handling? https://chalk-lab.github.io/Mooncake.jl/stable/understanding_mooncake/rule_system/ does not explain any of that.

[ChrisRackauckas-Claude]

Complete analysis: kwcall rules are present, the hang is elsewhere

1. The @from_rrule kwcall rules ARE registered and working

The @from_rrule(..., true) in DiffEqBaseMooncakeExt generates both positional and Core.kwcall variants of _is_primitive and rrule_wrapper. I verified all checks pass for both MinimalCtx and DefaultCtx:

# All return true:
Mooncake._is_primitive(DefaultCtx, ReverseMode, kw_sig_5arg)       # true
Mooncake._is_primitive(DefaultCtx, ReverseMode, kw_sig_mtk_types)  # true  
Mooncake._is_primitive(DefaultCtx, ReverseMode, kw_sig_full_kwargs) # true

The only signature that doesn't match is SensitivityADPassThrough (which is <: AbstractDEAlgorithm, not <: AbstractSensitivityAlgorithm), but that's only used for internal "don't-intercept" calls and is irrelevant here.

No doubling of rules for keyword arguments is needed. The true third argument to @from_rrule already handles it.

2. Non-MTK ODE problems work correctly

With Mooncake as the outer AD and InterpolatingAdjoint(autojacvec=MooncakeVJP()) as the sensealg, gradients match ForwardDiff:

const _sa = InterpolatingAdjoint(autojacvec = MooncakeVJP())
loss(p) = sum(solve(prob, Tsit5(); u0=[1.0], p=p, sensealg=_sa, saveat=0.1))
# build_rrule: 12s, value_and_pullback!!: 96s, gradient matches ForwardDiff ✓

The solve_up primitive is found and used correctly through the Core.kwcall dispatch.

3. The actual hang: LLVM chokes on MTK's enormous type parameters

For MTK problems, the hang occurs during LLVM compilation, not during rule lookup or abstract interpretation. The stack trace when killed shows:

CorrelatedValuePropagationPass → LazyValueInfo → getEdgeValue → (hangs)
  ↑ called from:
generate_ir → jl_compile_codeinst → jl_add_to_ee → LLVM optimizer

Root cause: Mooncake traces through DiffEqBase.solve to reach solve_up (the primitive). Inside solve, there are field accesses like prob.u0 and prob.p. For MTK-generated ODEProblems, the type parameters are enormous — the initialization_data field alone has ~2KB of nested type parameters (OverrideInitData{NonlinearProblem{..., GeneratedFunctionWrapper{..., RuntimeGeneratedFunction{...}}, ..., InitializationMetadata{ReconstructInitializeprob{...}, ...}, ...}}).

Mooncake generates tangent-handling code for these field accesses. Even though many sub-types have tangent_type = NoTangent overrides (from MTKMooncakeExt), the generated Julia IR is still enormous because the type parameters appear in method signatures and dispatch code. The LLVM optimizer can't process this IR in reasonable time.

4. Sub-operations that work vs. don't

Operation	Time	Result
SciMLStructures.replace(Tunable(), mtk_prob.p, pval)	19s	✅ Works
solve(plain_prob, Tsit5(); ...)	108s	✅ Works
solve(mtk_prob, Tsit5(); ...)	>300s	❌ LLVM hangs

The difference between plain prob and MTK prob is the type complexity of the problem struct that solve accesses.

5. Potential fixes

Short term (SciML side): Add more tangent_type overrides in MTKMooncakeExt. The OverrideInitData type is NOT currently covered:

Mooncake.tangent_type(::Type{<:SciMLBase.OverrideInitData}) = NoTangent

This would simplify tangent code for the ODEFunction's initialization_data field. Similar overrides may be needed for other complex SciML types.

Medium term (DiffEqBase side): Register DiffEqBase.solve (not just solve_up) as a Mooncake primitive. This would prevent Mooncake from ever tracing into solve's body, completely avoiding the MTK type issue:

# In DiffEqBaseMooncakeExt: make solve itself a primitive that delegates to solve_up rrule
@is_primitive MinimalCtx Tuple{typeof(DiffEqBase.solve), DiffEqBase.AbstractDEProblem, Vararg{Any}}

With a custom rrule!! that extracts u0/p/sensealg from kwargs and calls the existing solve_up rrule.

Long term (Mooncake side): Better handling of types with enormous type parameters — either by limiting type parameter expansion in generated IR, or by recognizing that NoTangent fields don't need their full type information in the generated code.

[yebai]

MTK seems close to the edge of what Julia currently supports.

I don't think there is any correctness issue here, but I agree that adding a few more custom tangent types and rules is the right trade-off.

[ChrisRackauckas]

No, the AI just isn't able to isolate it well. Figuring out this Mooncake bug is a bit difficult since I don't know the compiler well though.

[ChrisRackauckas]

I added overloads into SciMLBase and ModelingToolkit to cover the large type parameters, so that's all gone, but it's still an infinite loop hang in the Mooncake compiler. 🤷‍♂️ I'm pretty much out of leads here, and the AI is just guessing things that aren't true at this point, so someone would need to show me some of Mooncake's compiler for me to fix it myself at this point.

[AstitvaAggarwal]

Yeah I'm on it 👍

[AstitvaAggarwal]

Nice problem encountered here. #1115 has the fix (part 1) and the reason of why this happens, with SciML/ModelingToolkit.jl#4416 as the Part 2 PR.