Mooncake: EnsembleProblem tutorials still on Zygote

[ChrisRackauckas-Claude]

Three tutorials still use AutoZygote because Mooncake's rule compiler fails on EnsembleProblem.__solve (MooncakeRuleCompilationError / stack overflow):

• docs/src/examples/sde/optimization_sde.md (Example 1)
• docs/src/examples/sde/SDE_control.md
• docs/src/tutorials/data_parallel.md

The SDE_control tutorial also needs Zygote.@nograd CreateGrid replaced with Mooncake.@zero_adjoint Mooncake.DefaultCtx Tuple{typeof(CreateGrid), Any, Any}.

SciMLBase 3.2.0 added tmap/responsible_map Mooncake rrule!! for EnsembleProblem, which may resolve this once the ecosystem cascade (LinearSolve, DiffEqBase, OrdinaryDiffEq, DelayDiffEq, StochasticDiffEq) lands.

🤖 Generated with Claude Code

[wsmoses]

@ChrisRackauckas-Claude see if this works with Enzyme and if not open an issue with a MWE

[ChrisRackauckas-Claude]

@wsmoses I checked this with current Enzyme (v0.13.151). Enzyme does not work for differentiating through an EnsembleProblem solve, in either EnsembleSerial or EnsembleThreads mode.

All variants fail identically — plain Reverse, set_runtime_activity(Reverse), via Optimization AutoEnzyme(), and even with the per-trajectory inner sensealg pinned to EnzymeVJP():

Enzyme.Compiler.EnzymeNoDerivativeError:
No augmented forward pass found for jl_get_builtin_fptr
  [1] AggregateLogger @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:117
  [2] AggregateLogger @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:125
  [3] #__solve#733     @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:249

Importantly the per-trajectory adjoint is fine: differentiating the same loss with a Zygote outer pass and inner sensealg = InterpolatingAdjoint(autojacvec = EnzymeVJP()) gives the correct gradient. The failure is Enzyme tracing the ensemble orchestration, specifically the jl_get_builtin_fptr / Core._compute_sparams emitted when constructing SciMLBase's parametric AggregateLogger{T} inside __solve.

I reduced it to a SciML-free, Logging-free MWE — same missing rule as EnzymeAD/Enzyme.jl#2437 (jl_get_builtin_fptr):

import Enzyme
@noinline get_any()::Any = Base.inferencebarrier(Ref(1.0))
mutable struct Wrap{T}; v::Float64; inner::T; end
f(x) = (w = Wrap(0.0, get_any()); w.v = sum(abs2, x); w.v)   # T computed at runtime -> _compute_sparams
Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), f, [3.0, 1.0])
# EnzymeNoDerivativeError: No augmented forward pass found for jl_get_builtin_fptr

Because the trigger is the parametric AggregateLogger{T} construction, a SciML-side workaround may also be feasible (an Enzyme rule for the ensemble solve, or avoiding the runtime type-parameter computation in the logger wrap).

Verified on Julia 1.11.9 / Enzyme 0.13.151 / SciMLSensitivity 7.111.0 / OrdinaryDiffEq 7.0.0. (Note: current SciMLBase uses the 2-arg prob_func(prob, ctx) form — the tutorials' 3-arg (prob, i, repeat) form is itself stale and errors before AD even starts.)

[wsmoses]

@ChrisRackauckas-Claude that bug was resolved, please retry

[ChrisRackauckas-Claude]

@wsmoses Retried with Enzyme v0.13.154 (Julia 1.11.9, SciMLSensitivity v7.112.0, SciMLBase v3.18.0, OrdinaryDiffEq v7.0.0).

EnzymeAD/Enzyme.jl#2437 is indeed fixed — the No augmented forward pass found for jl_get_builtin_fptr error is gone. But the ensemble differentiation still does not work; the failure moved one layer deeper into the new code path.

1. set_runtime_activity(Reverse), both EnsembleSerial and EnsembleThreads: crashes in Enzyme's type-unstable newstruct shadow construction, at the same trigger site as before (AggregateLogger{ConsoleLogger} construction in basic_ensemble_solve.jl):

MethodError: no method matching getindex(::Nothing)
  [1] macro expansion          @ Enzyme/src/rules/typeunstablerules.jl:39 [inlined]
  [2] create_shadow_ret        @ Enzyme/src/rules/typeunstablerules.jl:3 [inlined]
  [4] runtime_newstruct_augfwd(::Type{Val{(true, false, false, false, true, false)}}, ::Val{true}, ::Val{1},
        ::Val{(true, true, true, true, true, true, true)}, ::Type{SciMLBase.AggregateLogger{Base.CoreLogging.ConsoleLogger}},
        ::Val{Any}, ::Dict{Symbol, Float64}, ::Dict{Symbol, Float64}, ::Int64, ::Nothing, ::Float64, ::Nothing,
        ::Float64, ::Nothing, ::ReentrantLock, ::ReentrantLock, ::Base.CoreLogging.ConsoleLogger, ::Nothing)
                               @ Enzyme/src/rules/typeunstablerules.jl:322
  [5] AggregateLogger          @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:117 [inlined]
  [7] #__solve#733             @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:249 [inlined]

The SciML-free MWE from my previous comment reproduces this new failure exactly (same create_shadow_ret stack), so it remains a valid MWE for the residual bug:

import Enzyme
@noinline get_any()::Any = Base.inferencebarrier(Ref(1.0))
mutable struct Wrap{T}; v::Float64; inner::T; end
f(x) = (w = Wrap(0.0, get_any()); w.v = sum(abs2, x); w.v)
Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), f, [3.0, 1.0])
# MethodError: no method matching getindex(::Nothing)
#   in create_shadow_ret / runtime_newstruct_augfwd, typeunstablerules.jl:39

2. Bypassing the logger with solve(...; progress_aggregate = false) (skips AggregateLogger construction entirely): Enzyme gets further into compilation but then kills the Julia process with native memory corruption — 3/3 runs crashed, at varying locations: twice signal 7 (Bus error) inside the LLVM pass pipeline (~PassBuilder / runJuliaPasses from compile_unhooked at Enzyme/src/compiler.jl:5356), once signal 11 (Segmentation fault) in GC (gc_mark_loop_serial) while allocating inside abs_typeof → check_ir! → nested_codegen! → enzyme_custom_augfwd, i.e. while processing a custom rule. Nondeterministic crash location during compilation smells like a GC rooting bug in the Enzyme compiler.

3. Plain Reverse (no runtime activity): EnzymeRuntimeActivityError as expected, so runtime activity is mandatory here, which lands on failure 1.

So: original missing-rule bug fixed, but EnsembleProblem adjoints via Enzyme remain blocked, now by (a) the getindex(::Nothing) bug in runtime_newstruct_augfwd shadow construction (MWE above), and behind that (b) a native crash during Enzyme compilation of the ensemble internals.

[ChrisRackauckas-Claude]

Follow-up: I dug into the residual failures and they turned out to be a chain of bugs. Three were Julia-side in Enzyme.jl and I have a working patch (verified locally against Enzyme main); the fourth is in Enzyme core C++ and still blocks the ensemble adjoint.

Bug 2 — getindex(::Nothing) in create_shadow_ret (typeunstablerules.jl:39). A constant argument (ActivityTup[i] == false) whose type is ActiveState — e.g. the Float64 literal initializing an AggregateLogger field — has shadow nothing, and create_shadow_ret has no case for it, falling through to batchshadowarg[]. Latent since the function was introduced (#2431); exposed once #2437's fix let compilation reach it. Fixed by giving constant+ActiveState args a make_zero(primarg) shadow, and gating the reverse-pass accumulation in body_construct_rev on ActivityTup[i] (it would otherwise throw EnzymeNonScalarReturnException on the same inputs). The MWE from my earlier comment exercises exactly this and passes with the fix (gradient [6.0, 2.0], correct).

Bug 3 — MethodError: Duplicated(::AggregateLogger{ConsoleLogger}, ::Base.RefValue{...}) in create_activity_wrapper. runtime_newstruct_augfwd Ref-wraps the shadow whenever any field is Active/Mixed state, but a mutable struct is DupState overall (active_reg_nothrow(AggregateLogger{ConsoleLogger}) == DupState), so downstream consumers derive "plain shadow object" from the type and choke on the RefValue. Fixed by Ref-wrapping only for immutable (genuinely MixedState) types, with the tape read in body_construct_rev mirrored.

Bug 4 — error reporter masking diagnostics: julia_error (errors.jl:1058) calls GPUCompiler.backtrace(val) on a ConstantExpr instead of its Instruction user → MethodError: backtrace(::LLVM.ConstantExpr). One-line fix.

With those fixed, test/typeunstable.jl (including the #2817 runtime-activity tuple tests), test/mixed.jl, and test/mixedapplyiter.jl all pass, plus a new regression test.

Bug 5 — remaining blocker, Enzyme core C++. The ensemble gradient now compiles much deeper and then kills the process inside Enzyme proper:

• Julia 1.10.11: deterministic SIGSEGV in llvm::Instruction::eraseFromParent via RecursivelyReplaceAddressSpace / UpgradeAllocasToMallocs (Enzyme/FunctionUtils.cpp:618/779) during CreateAugmentedPrimal → CreateFromClone → preprocessForClone of _dispatch_ensemble_solve.
• Julia 1.11.9: GC-corruption segfaults/bus errors at varying points mid-compilation (in the julia_allocator → zero_allocation callback invoked from C++ CreateAllocation; in abs_typeof during check_ir!; in the LLVM pass pipeline) — classic missing-GC-root signature.
• Right before the 1.11 crash, Enzyme emits an ET_GCRewrite warning — "Could not find use of stored value" for a store of @ejl_jl_nothing into augmented_julia__dispatch_ensemble_solve_*, whose sret is a 3584-byte aggregate with enzymejl_returnRoots="271". The GC-root rewriting appears to lose roots on this large sret/returnRoots aggregate, consistent with both crash modes.

Patch (against Enzyme main, with regression test) and full writeup: https://github.com/ChrisRackauckas/InternalJunk/pull/42. @wsmoses happy to send the typeunstablerules/errors patch as an Enzyme.jl PR if you want it (pending Chris's go-ahead); bug 5 has no reduced MWE yet — the trigger is the huge sret of _dispatch_ensemble_solve, so reduction would aim at a runtime-dispatched call returning a large many-GC-root aggregate.

Versions: Enzyme.jl main (post-v0.13.154, 14df74d0 base), SciMLSensitivity v7.112.0, SciMLBase v3.18.0, OrdinaryDiffEq v7.0.0, Julia 1.11.9 / 1.10.11.

[ChrisRackauckas-Claude]

The Enzyme.jl fixes are now up as PRs: EnzymeAD/Enzyme.jl#3162 (error-reporter fix) and EnzymeAD/Enzyme.jl#3163 (the type-unstable newstruct shadow fixes with regression tests). The remaining ensemble blocker is the Enzyme-core C++ crash described above.

[ChrisRackauckas-Claude]

Update: Enzyme v0.13.157 includes the reporter fix (EnzymeAD/Enzyme.jl#3162, merged) and @wsmoses's EnzymeAD/Enzyme.jl#3173 which subsumed #3163 — verified locally that both MWEs now pass on the release. The remaining native crash still reproduces on v0.13.157 (deterministic UpgradeAllocasToMallocs segfault on Julia 1.10.11, GC corruption abort on 1.11.9) and is now filed as EnzymeAD/Enzyme.jl#3175.

[wsmoses]

@ChrisRackauckas-Claude / @ChrisRackauckas can you retry with latest enzyme

[ChrisRackauckas-Claude]

@wsmoses Retried on the latest release, Enzyme v0.13.163 (SciMLSensitivity v7.112.0, SciMLBase v3.22.0, OrdinaryDiffEq v7.1.0, Julia 1.11.9 and 1.10.11). The EnsembleProblem adjoint is still blocked on both versions, but the failure mode on 1.10 has changed for the better — the process-killing native crash is gone there.

Confirmed fixed / working:

• The type-unstable newstruct shadow MWE passes on both 1.10.11 and 1.11.9 (gradient [6.0, 2.0]) — #3162/#3173 are in the release; bugs 2–4 stay fixed.
• The per-trajectory adjoint is fine and correct. A single (non-ensemble) solve with sensealg = InterpolatingAdjoint(autojacvec = EnzymeVJP()) differentiated by an outer Enzyme Reverse gives the right gradient on both versions, identical to a finite-difference reference ([45.741, -24.292, 6.161, -7.638]). So the blocker is purely the ensemble orchestration.

Still failing — Enzyme.gradient(set_runtime_activity(Reverse), loss, p) where loss does solve(EnsembleProblem(...), Tsit5(), EnsembleSerial(); trajectories=3, sensealg=InterpolatingAdjoint(autojacvec=EnzymeVJP())):

• Julia 1.11.9 — unchanged, still #3175. SIGSEGV from GC corruption mid-compilation (gc_mark_loop_serial → jl_datatype_layout, triggered by a Box alloc in check_ir! reached via nested_codegen! → enzyme_custom_augfwd, inside a deeply-recursive CreateAugmentedPrimal over the ensemble internals). Immediately before the crash Enzyme emits the same ET_GCRewrite warning "Could not find use of stored value" for a store of @ejl_jl_nothing into augmented_julia__dispatch_ensemble_solve_*_inner_*wrap, whose sret is a dereferenceable(3576) aggregate with enzymejl_returnRoots="264". Same large-sret/returnRoots GC-root-rewriting bug as before.

• Julia 1.10.11 — the native crash is GONE. The deterministic UpgradeAllocasToMallocs SIGSEGV no longer reproduces; compilation now runs all the way through and fails with a clean, recoverable Julia-level error instead:

  AssertionError: Unimplemented accumulate_into for type Tuple{Float64, Float64}
    @ Enzyme src/internal_rules/core.jl:194
    ...
    runtime_generic_rev(... ::typeof(deepcopy), ::ODEProblem{...}, ...)
    batch_func @ SciMLBase/src/ensemble/basic_ensemble_solve.jl:396
    responsible_map @ .../basic_ensemble_solve.jl:491
  

  The trigger is differentiating the deepcopy(prob) (the ensemble safetycopy) inside batch_func, which needs to accumulate_into the ODE's tspan::Tuple{Float64,Float64} — and there's no accumulate_into method for tuples.

Workaround attempt (negative result): setting safetycopy = false skips that deepcopy, which avoids the assertion on 1.10.11 — but the gradient that comes back is numerically garbage (≈ 1e45–1e46, vs. the correct O(10)), so it is not a usable workaround.

Net: nothing regressed, the earlier Julia-side fixes hold, and the per-trajectory path is correct — but the ensemble adjoint is still blocked. On 1.11.9 it's the #3175 GC-root-rewrite crash on the huge _dispatch_ensemble_solve sret; on 1.10.11 that crash is now downgraded to a missing accumulate_into(::Tuple) rule on the deepcopy-of-ODEProblem path. Happy to file the accumulate_into(::Tuple) gap separately if useful.

Reproduce: loss(p) = (eprob = EnsembleProblem(remake(prob, p=p), prob_func = (pr,ctx)->remake(pr, u0 = pr.u0 .+ 0.05*ctx.sim_id)); sol = solve(eprob, Tsit5(), EnsembleSerial(); trajectories=3, sensealg=InterpolatingAdjoint(autojacvec=EnzymeVJP())); sum(sum(abs2,Array(s)) for s in sol.u)) with a Lotka–Volterra prob, then Enzyme.gradient(set_runtime_activity(Reverse), loss, p).