Mooncake: second-order adjoints (forward-over-reverse Hessian) not yet working

[ChrisRackauckas-Claude]

docs/src/examples/ode/second_order_adjoints.md uses Mooncake for Phase 1 (Adam) but still uses Zygote for Phase 2 (NewtonTrustRegion) because Mooncake has no working forward-over-reverse path through SciMLSensitivity + OrdinaryDiffEq.

• SecondOrder(AutoMooncake(), AutoMooncake()) raises "reverse-over-reverse not supported" (ArgumentError from build_derived_rrule).
• SecondOrder(AutoForwardDiff(), AutoMooncake()) is blocked on Mooncake's IEEEFloat-only gradient interface — it rejects ForwardDiff.Dual as the primal type (ValueAndGradientReturnTypeError).

Mooncake's native prepare_hvp_cache / value_and_hvp!! works on simple functions (Rosenbrock), but on the neural ODE loss through OrdinaryDiffEq it hits:

1. Missing foreigncall frules for :jl_get_world_counter and :jl_matching_methods — fixed in Add frule!!/rrule!! for :jl_get_world_counter and :jl_matching_methods chalk-lab/Mooncake.jl#1142.
2. UndefRefError in _build_output_tangent_cartesian for IdDict tangent construction (Memory{Any} has #undef slots) — separate Mooncake bug.
3. ArgumentError: Tangent types do not match primal types for ComponentArray closures captured into HVPCache's grad_f — separate Mooncake/CA issue.

Once Mooncake's forward-over-reverse matures, both phases can use Mooncake.

🤖 Generated with Claude Code

[wsmoses]

@ChrisRackauckas-Claude see if this works with enzyme

[ChrisRackauckas-Claude]

@wsmoses checked with current Enzyme (v0.13.151). Enzyme does not work for this second-order / forward-over-reverse Hessian path either, and it fails earlier than Mooncake.

Reference path that works at every level: SecondOrder(AutoForwardDiff(), AutoZygote()) — Hessian (252×252), HVP, and the full NewtonTrustRegion solve (objective 189.1 → 98.6). So the forward-over-reverse approach and the MWE are sound; only the Enzyme inner reverse fails.

Every Enzyme path fails:

• SecondOrder(AutoForwardDiff(), AutoEnzyme(set_runtime_activity(Reverse))) (the key combo) → MethodError: getindex(::Nothing) in create_shadow_ret / runtime_newstruct_augfwd — the inner reverse already fails (building the @NamedTuple{saveat} kwarg of the ODE solve), so second order is never reached. Same as Cannot store inactive but differentiable variable into active tuple EnzymeAD/Enzyme.jl#2212.
• plain AutoEnzyme() → EnzymeRuntimeActivityError.
• SecondOrder(AutoEnzyme(), AutoEnzyme()) → EnzymeNoDerivativeError: No forward mode derivative found for jl_field_isdefined_checked in Enzyme make_zero.

Recommendation: keep Phase 2 on Zygote — SecondOrder(AutoForwardDiff(), AutoZygote()) is the only forward-over-reverse path that currently works through SciMLSensitivity + OrdinaryDiffEq. The issue's premise extends to Enzyme: Enzyme forward-over-reverse through the neural ODE doesn't work yet, failing in the inner reverse on the same kwarg-NamedTuple bug. Verified on Julia 1.11.9 / Enzyme 0.13.151 / SciMLSensitivity 7.111.0 / OrdinaryDiffEq 7.0.0.

[ChrisRackauckas-Claude]

Correction / refinement to my comment above (prompted by the observation that #1427 is fundamentally about the AD nesting — so I re-tested with that isolated).

My earlier comment blamed the failure on Bug A (the getindex(::Nothing) kwarg-NamedTuple issue, EnzymeAD/Enzyme.jl#2212) firing in the inner reverse, implying the second-order nesting was "never reached." That diagnosis was incomplete. Once Bug A is removed — making the solve kwarg-source globals (saveat = tsteps) const — the picture changes:

First-order Enzyme reverse actually works here, with the default (auto-selected) sensealg: Enzyme.gradient(set_runtime_activity(Reverse), loss_neuralode, ps) is finite and matches Zygote to relerr 0.0. So the inner reverse is fine and my "never reached" claim was wrong. (Caveat: forcing an explicit sensealg as a const gives numerically wrong Enzyme gradients — Enzyme then differentiates through the discrete-adjoint machinery — so only the default sensealg is correct.)

The genuine forward-over-reverse nesting still fails, but at a new site (not Bug A). SecondOrder(AutoEnzyme(mode = Enzyme.Forward), AutoEnzyme(mode = set_runtime_activity(Enzyme.Reverse))) — Enzyme forward outer over Enzyme reverse inner — throws:

EnzymeNoDerivativeError: No forward mode derivative found for jl_field_isdefined_checked
  [1] make_zero @ Enzyme/src/typeutils/make_zero.jl:255

i.e. the outer forward pass can't differentiate the make_zero shadow-allocator the inner reverse uses (the in-place Optimization/NewtonTrustRegion path fails the same way via jl_idset_peek_bp in make_zero!, make_zero!.jl:583). SecondOrder(AutoForwardDiff(), AutoEnzyme(Reverse)) fails differently again — the inner Enzyme reverse can't return a ForwardDiff.Dual ("Active return values … not type ForwardDiff.Dual"). I did not find an existing Enzyme.jl issue for this make_zero forward-over-reverse gap.

Separately, even the (first-call-correct) first-order Enzyme reverse appears to corrupt persistent solver/library state: it's exact on the first call, but a subsequent plain forward loss_neuralode(ps) (with ps const/unchanged) jumps from 125.7 to 3.8e4, and repeated gradients — Enzyme and Zygote — then explode by orders of magnitude. So Enzyme reverse isn't reliably usable here even at first order.

Net: the recommendation is unchanged — keep Phase 2 on SecondOrder(AutoForwardDiff(), AutoZygote()) (verified: drives NewtonTrustRegion, objective decreases). But the reason is the make_zero forward-over-reverse gap, not Bug A; and first-order Enzyme reverse should no longer be described as outright failing (it works once, modulo the state-corruption caveat). Verified on Julia 1.11.9 / Enzyme 0.13.151 / SciMLSensitivity 7.111.0 / OrdinaryDiffEq 7.0.0.

---

Update: filed the make_zero forward-over-reverse gap upstream as EnzymeAD/Enzyme.jl#3135 (reduced to a SciML-free MWE: nested Enzyme forward-over-reverse where the inner reverse's make_zero recurses into an undef-able/abstract field).

[wsmoses]

@ChrisRackauckas-Claude / @ChrisRackauckas can you retest with enzyme for the current status [opening issues for any remainders]

[ChrisRackauckas-Claude]

@wsmoses retested forward-over-reverse for the neural-ODE loss on the current stack: Enzyme 0.13.163, SciMLSensitivity 7.112.0, OrdinaryDiffEq 7.1.0, DifferentiationInterface 0.7.18, Julia 1.11.9. (Previous retest was Enzyme 0.13.151 / SciMLSensitivity 7.111.0 / OrdinaryDiffEq 7.0.0.) The MWE is the docs second_order_adjoints.md loss (252-param Lux neural ODE, ComponentArray p).

Summary

path	result
first-order Enzyme.gradient(set_runtime_activity(Reverse), loss, ps) (default sensealg)	✅ correct — relerr 0.0 vs Zygote… but corrupts state (see below)
SecondOrder(AutoForwardDiff(), AutoEnzyme(Reverse)) — HVP / Hessian / NewtonTrustRegion	❌ inner Enzyme reverse can't return a ForwardDiff.Dual active value
SecondOrder(AutoEnzyme(Forward), AutoEnzyme(Reverse)) — HVP	❌ Enzyme: Not yet implemented forward for jl_eqtable_get in make_zero (the #3135 cascade)
reference SecondOrder(AutoForwardDiff(), AutoZygote()) — HVP / Hessian (252×252) / NewtonTrustRegion	✅ works

Net recommendation is unchanged: keep Phase 2 on SecondOrder(AutoForwardDiff(), AutoZygote()) — still the only forward-over-reverse path that works end-to-end through SciMLSensitivity + OrdinaryDiffEq. But the status of the two Enzyme failure modes has moved, and there's a new first-order finding.

1. Forward-over-reverse make_zero cascade has advanced (good news)

Last time, SecondOrder(AutoEnzyme(Forward), AutoEnzyme(Reverse)) failed with EnzymeNoDerivativeError: No forward mode derivative found for jl_field_isdefined_checked (filed as EnzymeAD/Enzyme.jl#3135, fixed by EnzymeAD/Enzyme.jl#3137 — merged). On 0.13.163 that error is gone; the nesting now gets one builtin further into make_zero's IdDict seen-table bookkeeping and stops at:

EnzymeRuntimeException: Enzyme: Not yet implemented forward for jl_eqtable_get
  [4] make_zero  @ Enzyme/src/typeutils/make_zero.jl:39   (haskey(seen, prev))
  [5] make_zero  @ Enzyme/src/typeutils/make_zero.jl:257
  ...
  [8] gradient   @ Enzyme/src/sugar.jl:274                 (inner reverse shadow init)
 [21] hvp ...

This is exactly the jl_eqtable_* follow-up cascade predicted in EnzymeAD/Enzyme.jl#3135. It reproduces on the released Enzyme both through the full neural ODE and via the reduced Enzyme-only MWE from that issue (control with a plain Vector{Float64} and no undef-able field still returns the correct HVP). Updating #3135 with the on-release status.

2. SecondOrder(AutoForwardDiff(), AutoEnzyme(Reverse)) — still blocked, unchanged

Active return values with automatic pullback (differential return value) deduction only supported
for floating-like values and not type ForwardDiff.Dual{...}. ... you can explicitly specify a
pullback by using split mode, e.g. autodiff_thunk(ReverseSplitWithPrimal, ...)

The inner Enzyme reverse refuses a ForwardDiff.Dual active return — so ForwardDiff-outer-over-Enzyme-inner can't form an HVP/Hessian. Same failure for DI.hvp, DI.hessian, and the Optimization/NewtonTrustRegion in-place path. (DI's Enzyme reverse extension would need split mode to support a Dual-eltype HVP.)

3. NEW: first-order Enzyme reverse silently mutates the captured u0 (state corruption)

The first-order Enzyme reverse gradient is numerically correct (relerr 0.0 vs Zygote) on the first call, but the call mutates the captured u0 initial-condition array (the primal). Everything afterward is poisoned:

[1] plain forward loss(ps)              = 179.87581
[4] Enzyme reverse gradient   relerr    = 0.0          # correct
[5] plain forward loss(ps) AFTER Enzyme = 322705.75    # u0 has been clobbered
[7] Zygote gradient AFTER Enzyme  relerr = 4327.6      # everything downstream explodes

Isolating what changed: ps unchanged, ode_data unchanged, u0 changed (maxabsdiff ≈ 101). Restoring u0 recovers the original loss exactly, and passing a fresh copy(u0) into the ODEProblem per call eliminates the corruption entirely (loss stable, u0 stays [2.0, 0.0], repeated gradients all relerr 0.0). So the reverse rule for solve is writing the u0 cotangent into the u0 primal buffer rather than into a separate shadow. This is a first-order bug, independent of the second-order question — filed separately as #1494. (Possibly related to #1469 / #1470, but this is the minimal non-MTK neural-ODE case.)

All four MWEs (full retest harness, isolated corruption repro, copy(u0) workaround, reduced Enzyme-only forward-over-reverse) are reproducible; happy to attach.

[wsmoses]

File dependency free mwes for all issues

[wsmoses]

@ChrisRackauckas-Claude the deepcopy issue should be fixed, retry. Note that ForwardDiff cannot differentiate through Enzyme code so that failure is expected [and should be a louder error in forwarddiff]