ext/Mooncake: bridge copyto!(::ComponentArray, ::Mooncake.Tangent)

[ChrisRackauckas-Claude]

Please ignore until reviewed by @ChrisRackauckas.

Summary

Bridge copyto!(::ComponentArray, ::Mooncake.Tangent) for both the flat-Array-backed and SubArray-backed ComponentVector tangent shapes so Mooncake gradients can be written back into a ComponentArray-shaped buffer.

The bug

using Optimization, OptimizationOptimJL, ComponentArrays
import Mooncake

f(x, _) = sum(abs2, x)
optf = OptimizationFunction(f, AutoMooncake())

# works
solve(OptimizationProblem(optf, randn(100)), LBFGS(); maxiters=10)

# fails:
#   MethodError: no method matching iterate(
#     ::Mooncake.Tangent{@NamedTuple{data::Vector{Float64}, axes::Mooncake.NoTangent}})
solve(OptimizationProblem(optf, ComponentArray(a=randn(100))), LBFGS(); maxiters=10)

Root cause

DifferentiationInterface.value_and_gradient!(::AutoMooncake, …) writes the gradient into the user-supplied grad buffer with an unconditional copyto!(grad, new_grad) (DI's DifferentiationInterfaceMooncakeExt/onearg.jl:153). For a plain Vector{Float64} primal this works because Mooncake.tangent_type(Vector{Float64}) === Vector{Float64}, so the source is an array. For a ComponentArray primal, however, Mooncake.tangent_type is a Mooncake.Tangent struct — not an AbstractArray — so the generic Base.copyto!(::AbstractArray, ::Any) fallback drops into the iterate-based path at abstractarray.jl:934 and throws MethodError for iterate.

DI's twoarg and forward call sites already handle the "tangent ≠ primal-shape" case via dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy). The value_and_gradient! path is the one missed case; a parallel issue could be filed against DI to apply the same pattern there. Until that lands, defining copyto! for the (ComponentArray, Mooncake.Tangent) pairs fully unblocks the user-visible failure in Optimization.jl + AutoMooncake + ComponentArrays.

Coverage

I probed the obvious variations against this branch. Results:

#	Case	Status
1	Flat 1D ComponentVector	✅ original failing case — fixed by overload (a)
2	Nested CV (ComponentArray(u=…, p=ComponentArray(…)))	✅ still flat-backed underneath; same overload (a)
3	ComponentMatrix (2D backing)	✅ Matrix{P} <: Array{P} — overload (a) covers it; regression test added
4	Float32 CV	✅ _FloatLike covers it
5	SubArray-backed CV, full cover	✅ new overload (b); test added
6	SubArray-backed CV, partial cover	✅ clear ArgumentError (symmetric to existing _increment_subarray_fdata! guard); test added
7	Non-mutating DI.gradient(f, prep, ::AutoMooncake, ::CV)	⚠️ returns a Mooncake.Tangent, not a CV — outside CA's reach, needs DI/Mooncake-level fix
8	OptimizationProblem(optf, parent_cv.u) where u is a flat 1D component	⚠️ parent_cv.u returns a raw SubArray, not a CV — Mooncake-on-SubArrays issue, outside CA's scope

(7) is a latent DI contract issue: the non-mutating DI.gradient path runs _copy_output(new_grad) and returns the deep-copied Tangent, so the caller gets back a type that doesn't match the input's shape. The right fix is in DI (mirror the _copy_to_output!! pattern) or Mooncake's _copy_output dispatch. Mentioning here for visibility but deliberately not addressing in this PR.

(8) requires a Mooncake-side copyto!(::Vector, ::Mooncake.Tangent{…SubArray-tangent…}) since the destination isn't a ComponentArray at all once the wrapper is dropped.

What this PR does

• ext/ComponentArraysMooncakeExt.jl: adds two Base.copyto! overloads —
  • (a) (ComponentArray{P,N,<:Array{P}}, Tangent{@NamedTuple{data::Array{P}, axes::NoTangent}}) → copies src.fields.data into getdata(dest).
  • (b) (ComponentArray{P,N,<:AbstractArray{P}}, Tangent{@NamedTuple{data::Tangent{@NamedTuple{parent::Vector{P}, indices::NoTangent, offset1::NoTangent, stride1::NoTangent}}, axes::NoTangent}}) → copies src.fields.data.fields.parent into getdata(dest), guarded by a length(parent) == length(getdata(dest)) full-cover check that throws a clear ArgumentError on partial-cover views (symmetric to the existing _increment_subarray_fdata! block).
• test/autodiff/autodiff_tests.jl: adds four targeted tests in the Mooncake testset — flat CV, ComponentMatrix, full-cover SubArray-backed CV, and @test_throws ArgumentError for partial-cover SubArray-backed CV.
• Project.toml: bumps version = "0.15.39".

Local verification

Autodiff test group passes with Mooncake testset at 15/15 (was 7/7 on main).

End-to-end with the dev'd ComponentArrays:

---- ComponentArray case ----
retcode = Success   objective = 0.0
typeof(sol.u) = ComponentVector{Float64, Vector{Float64}, …}
OK

Previous CI on the (a)-only commit was 18/18 green (all matrix entries for Julia 1, lts, pre × Core/Autodiff/Downstream/Reactant/GPU/nopre, plus Runic, build, doc-build).

Test plan

• CI: Core group passes
• CI: Autodiff group passes (new copyto!(::CV, ::Mooncake.Tangent) cases in test/autodiff/autodiff_tests.jl)
• CI: Downstream group passes
• CI: Reactant group passes
• CI: nopre group passes

[gdalle]

I'm not sure that this is the correct fix. I think since Mooncake v0.5.25 (more specifically since chalk-lab/Mooncake.jl#1103), every package that defines types with "friendly tangents" needs to add some overloads. See e.g. the similar discussion around StaticArrays:

• Add Mooncake integration JuliaArrays/StaticArrays.jl#1342
• Add Mooncake.jl extension JuliaArrays/StaticArrays.jl#1343

as well as more general background:

• Changes in Mooncake 0.5.25 that may affect DI users JuliaDiff/DifferentiationInterface.jl#986
• StaticArrays need opt-in friendly tangent type for Mooncake JuliaDiff/DifferentiationInterface.jl#998
• fix: make DI compatible with latest Mooncake friendly tangents JuliaDiff/DifferentiationInterface.jl#1001

[ChrisRackauckas]

It probably could be handled more generically, but with the current development progression I have been suggested to go this direction, and this patch at least unblocks users for today. If things change in the core repo then we can remove the patch.