Add Mooncake extension for ArrayPartition cotangents

When an upstream ChainRules-based adjoint (e.g. SciMLSensitivity's
`_concrete_solve_adjoint` for an ODE whose state is an `ArrayPartition`,
such as the one produced by `SecondOrderODEProblem`) returns a
parameter / state cotangent as an `ArrayPartition`, Mooncake's
`@from_chainrules` / `@from_rrule` accumulator looks for an
`increment_and_get_rdata!` method matching

    (FData{NamedTuple{(:x,), Tuple{Tuple{Vector, …}}}}, NoRData, ArrayPartition)

There isn't a default method registered for this combination, so the
call falls through to the generic error path:

    ArgumentError: The fdata type Mooncake.FData{@NamedTuple{x::Tuple{Vector{Float32}, Vector{Float32}}}},
    rdata type Mooncake.NoRData, and tangent type
    RecursiveArrayTools.ArrayPartition{Float32, Tuple{Vector{Float32}, Vector{Float32}}}
    combination is not supported with @from_chainrules or @from_rrule.

Add the missing dispatch via a new `RecursiveArrayToolsMooncakeExt`
weak-dep extension. An `ArrayPartition`'s only field is `x::Tuple` of
inner arrays, so the FData layout is `FData{@NamedTuple{x::Tuple{...}}}`
and the inner tuple positions line up with `t.x`. Walk the tuple
element-by-element and forward each leaf to the existing
`increment_and_get_rdata!` for the leaf's array type, which does the
actual in-place accumulation. Returns `Mooncake.NoRData()` to match the
no-rdata convention used by the equivalent ComponentArrays dispatch
(SciML/ComponentArrays.jl#350 / #351).

Tested end-to-end against the SciMLSensitivity neural-ODE
`SecondOrderODEProblem` tutorial (via SciML/SciMLSensitivity.jl#1422,
which adds the matching `df_iip`/`df_oop` cotangent unwrap on the
SciMLSensitivity side): with both PRs applied, the Lux + `ArrayPartition`
training loop now runs under
`OPT.AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true))`.

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

Author: ChrisRackauckas

Project.toml

@@ -20,6 +20,7 @@ FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -37,6 +38,7 @@ RecursiveArrayToolsFastBroadcastPolyesterExt = ["FastBroadcast", "Polyester"]
 RecursiveArrayToolsForwardDiffExt = "ForwardDiff"
 RecursiveArrayToolsKernelAbstractionsExt = "KernelAbstractions"
 RecursiveArrayToolsMeasurementsExt = "Measurements"
+RecursiveArrayToolsMooncakeExt = "Mooncake"
 RecursiveArrayToolsMonteCarloMeasurementsExt = "MonteCarloMeasurements"
 RecursiveArrayToolsReverseDiffExt = ["ReverseDiff", "Zygote"]
 RecursiveArrayToolsSparseArraysExt = ["SparseArrays"]
@@ -58,6 +60,7 @@ GPUArraysCore = "0.2"
 KernelAbstractions = "0.9.36"
 LinearAlgebra = "1.10"
 Measurements = "2.11"
+Mooncake = "0.5"
 MonteCarloMeasurements = "1.2"
 NLsolve = "4.5"
 Pkg = "1"

ext/RecursiveArrayToolsMooncakeExt.jl

@@ -0,0 +1,44 @@
+module RecursiveArrayToolsMooncakeExt
+
+using RecursiveArrayTools
+using Mooncake
+
+# `ArrayPartition` cotangent handling for `@from_chainrules` /
+# `@from_rrule`-generated rules.
+#
+# When an upstream ChainRules-based adjoint (e.g. SciMLSensitivity's
+# `_concrete_solve_adjoint` for an ODE whose state is an `ArrayPartition`,
+# such as the one produced by `SecondOrderODEProblem`) returns a parameter
+# / state cotangent as an `ArrayPartition`, Mooncake's
+# `@from_chainrules`/`@from_rrule` accumulator looks for an
+# `increment_and_get_rdata!` method matching
+# `(FData{NamedTuple{(:x,), Tuple{Tuple{Vector, …}}}}, NoRData, ArrayPartition)`
+# — and there isn't one by default, so the call falls through to the
+# generic error path:
+#
+#     ArgumentError: The fdata type Mooncake.FData{@NamedTuple{x::Tuple{Vector{Float32}, Vector{Float32}}}},
+#     rdata type Mooncake.NoRData, and tangent type
+#     RecursiveArrayTools.ArrayPartition{Float32, Tuple{Vector{Float32}, Vector{Float32}}}
+#     combination is not supported with @from_chainrules or @from_rrule.
+#
+# Add the missing dispatch.  An `ArrayPartition`'s only field is `x::Tuple`
+# of inner arrays, so the FData layout is
+# `FData{@NamedTuple{x::Tuple{...}}}` and the inner tuple positions line up
+# with `t.x`.  Walk the tuple element-by-element and forward each leaf to
+# the existing `increment_and_get_rdata!` for the leaf's array type, which
+# does the actual in-place accumulation.
+function Mooncake.increment_and_get_rdata!(
+        f::Mooncake.FData{@NamedTuple{x::T}},
+        r::Mooncake.NoRData,
+        t::ArrayPartition{P, T},
+    ) where {P, T <: Tuple}
+    fxs = f.data[:x]
+    txs = t.x
+    @assert length(fxs) == length(txs)
+    for i in eachindex(fxs)
+        Mooncake.increment_and_get_rdata!(fxs[i], Mooncake.NoRData(), txs[i])
+    end
+    return Mooncake.NoRData()
+end
+
+end