Remaining Mooncake SymbolicIndexingInterface AD gaps

After #1314, Mooncake supports the basic `sol[sym]` and `sol[sym, j]` paths on `AbstractTimeseriesSolution` for both state variables and scalar observables, including parameter gradients through observables. This issue tracks the remaining gaps for full feature parity with Zygote on the SymbolicIndexingInterface AD path. Tracks remaining items from #1207. ## What works after #1314 ```julia using ModelingToolkit, OrdinaryDiffEq, Mooncake using ModelingToolkit: t_nounits as t, D_nounits as D using DifferentiationInterface, ADTypes @parameters σ ρ β @variables x(t) y(t) z(t) w(t) eqs = [D(D(x)) ~ σ*(y-x), D(y) ~ x*(ρ-z) - y, D(z) ~ x*y - β*z, w ~ x+y+z+2β] @mtkcompile sys = System(eqs, t) prob = ODEProblem(sys, [...], (0.0, 100.0)) sol = solve(prob, Tsit5()) # ✓ State indexing gs = DifferentiationInterface.gradient(sol -> sum(sol[sys.x]), AutoMooncake(), sol) # ✓ Observable indexing (state + parameter gradient) gs = DifferentiationInterface.gradient(sol -> sum(sol[sys.w]), AutoMooncake(), sol) # gs.fields.u[k] = [1, 1, 1, 0] per timestep # gs.fields.prob.fields.p.fields.tunable[3] = 2 * length(sol) (∂(sum w)/∂β) ``` ## Remaining gaps ### 1. NonlinearSolution observable indexing (`isol[w]`) ```julia iprob = prob.f.initialization_data.initializeprob isol = solve(iprob) gs = DifferentiationInterface.gradient(isol -> isol[w], AutoMooncake(), isol) # ERROR: TypeError: non-boolean (Num) used in boolean context ``` **Root cause:** Inside MTK's `ObservedFunctionCache` for a `NonlinearSolution`, the dispatch path performs a symbolic comparison whose `==` overload returns a `Num`, not a `Bool`. Mooncake's source-to-source tracer can't differentiate through this. Zygote's operator overloading approach handles it because Zygote traces values, not source code. **Verified workarounds (not yet integrated):** - ForwardDiff handles this pattern. A combined `Mooncake + ForwardDiff` extension that uses ForwardDiff for the inner observed-function gradient when called from a Mooncake rrule works in principle. - Upstream fix in MTK: replace the `==` in `ObservedFunctionCache` lookup with an `===` or `objectid`-based comparison that doesn't trigger `Symbolics.==`. ### 2. Vector observable indexing (`sol[[sys.w, sys.x]]`) ```julia gs = DifferentiationInterface.gradient( sol -> sum(sum.(sol[[sys.w, sys.x]])), AutoMooncake(), sol ) # Falls back to differentiating through the [...] AbstractArray constructor, # which hits the SII dispatch chain (hash-consed symbols, atomic ops, etc.) ``` **Note:** This is also `@test_broken` for Zygote on Julia 1.10 (returns incorrect/swapped gradient elements). Tracked in #1233. Not a Mooncake-specific limitation. **Fix approach:** Add a Mooncake `rrule!!` for `getindex(::AbstractTimeseriesSolution, ::AbstractVector)` that decomposes the vector index into individual scalar `getindex` calls and accumulates the per-symbol pullbacks. ### 3. DAE observable indexing (`sol[sys.ampermeter.i]`) The Mooncake `getindex` rrule handles DAE solutions structurally (they're `AbstractTimeseriesSolution`), and the rrule does run. However, the existing test (`test/downstream/observables_autodiff.jl`) uses hard-coded expected values `[0.2, 1.0]` that don't match what either Zygote or Mooncake produce on the current test environment — the Zygote variant fails too. Likely stale expected values from an older MTK build, unrelated to Mooncake. **Fix:** Refresh the DAE test expected values, then enable the Mooncake test. ### 4. Full DAE adjoint with parameter gradient ```julia function loss_wrt_tunables(new_tunables) new_p = SS.replace(SS.Tunable(), prob.p, new_tunables) new_prob = remake(prob, p = new_p) sol = solve(new_prob, Rodas4()) return sum(sol[sys.ampermeter.i]) end ``` `@test_broken` for both Zygote and Mooncake. Pre-existing failure tracked in #1233 — `ChainRules` issue with `ModelingToolkitBase.PConstructorApplicator` (`Tuple field type cannot be Union{}`). Not specific to Mooncake. ### 5. Initialization observable parameter-gradient comparison The Zygote test does: ```julia @test gs.prob.p == gp ``` where `gp` is computed via `gradient(p -> f(state, p), backend, p)`. For Mooncake we'd need both #1 above (NonlinearSolution `isol[w]` working) and a way to compare Mooncake's `Tangent` parameter gradient with the Zygote-style NamedTuple from the inner gradient call. The accessor and comparison helpers would need a small abstraction. ## Suggested priority 1. **#3** is a quick win — refresh test expected values, both backends benefit. 2. **#1** is the highest-impact remaining gap. The cleanest path is the upstream MTK fix; the workaround via combined extension is also feasible. 3. **#2** depends on #1233 being resolved for Zygote first. 4. **#4** depends on #1233. 🤖 Generated with [Claude Code](https://claude.com/claude-code)