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

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <47193959+jonniedie@users.noreply.github.com>"]
-version = "0.15.38"
+version = "0.15.39"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/ComponentArraysMooncakeExt.jl

@@ -112,4 +112,61 @@ function Mooncake.friendly_tangent_cache(x::ComponentArray)
     return Mooncake.FriendlyTangentCache{Mooncake.AsPrimal}(copy(x))
 end
 
+# === Tangent → ComponentArray gradient copy ===========================================
+# `DifferentiationInterface.value_and_gradient!(::AutoMooncake, …)` writes the gradient
+# into a user-supplied `ComponentArray` buffer with an unconditional
+# `copyto!(grad, new_grad)`. Mooncake's `tangent_type` for a `ComponentArray` is a
+# `Mooncake.Tangent` struct, which is not an `AbstractArray` — so the generic
+# `Base.copyto!(::AbstractArray, ::Any)` fallback tries to iterate the tangent and
+# fails with a `MethodError` for `iterate`. Bridge both Tangent shapes that arise.
+
+# (a) Flat-Array-backed CV: tangent_type is
+#     `Tangent{@NamedTuple{data::Vector{P}, axes::NoTangent}}`.
+function Base.copyto!(
+        dest::ComponentArray{P, N, <:Array{P}},
+        src::Mooncake.Tangent{@NamedTuple{data::A, axes::Mooncake.NoTangent}},
+    ) where {P <: _FloatLike, N, A <: Array{P}}
+    copyto!(getdata(dest), src.fields.data)
+    return dest
+end
+
+# (b) SubArray-backed CV (from `getproperty(::ComponentVector, ::Symbol)` on a nested
+#     parent): tangent_type nests an inner Tangent that mirrors the SubArray's fields.
+#     Symmetric to the `_increment_subarray_fdata!` path already in this file: copy is
+#     only well-defined when the view fully covers its parent, since the SubArray
+#     indices are not recoverable from Mooncake fdata/tangent shape alone.
+function Base.copyto!(
+        dest::ComponentArray{P, N, <:AbstractArray{P}},
+        src::Mooncake.Tangent{
+            @NamedTuple{
+                data::Mooncake.Tangent{
+                    @NamedTuple{
+                        parent::Array{P, 1},
+                        indices::Mooncake.NoTangent,
+                        offset1::Mooncake.NoTangent,
+                        stride1::Mooncake.NoTangent,
+                    },
+                },
+                axes::Mooncake.NoTangent,
+            },
+        },
+    ) where {P <: _FloatLike, N}
+    parent = src.fields.data.fields.parent
+    if length(parent) != length(getdata(dest))
+        throw(
+            ArgumentError(
+                "ComponentArraysMooncakeExt: cannot copy a SubArray-backed " *
+                    "ComponentVector tangent (parent length $(length(parent))) into a " *
+                    "ComponentArray destination of length $(length(getdata(dest))). This " *
+                    "happens when a tangent flows out of a view that does not fully cover " *
+                    "its parent; there is no way to recover the view indices from Mooncake " *
+                    "tangent fields alone. Please file an issue against ComponentArrays.jl " *
+                    "with a reproducer.",
+            ),
+        )
+    end
+    copyto!(getdata(dest), parent)
+    return dest
+end
+
 end

test/autodiff/autodiff_tests.jl

@@ -204,4 +204,70 @@ end
 
     @test Mooncake.friendly_tangent_cache(flat) isa
         Mooncake.FriendlyTangentCache{Mooncake.AsPrimal}
+
+    # `copyto!(::ComponentVector, ::Mooncake.Tangent)` — required so that
+    # `DifferentiationInterface.value_and_gradient!(::AutoMooncake, …)` can write a
+    # Mooncake gradient back into a ComponentArray-shaped `grad` buffer. Without this
+    # bridge the generic AbstractArray `copyto!` fallback tries to iterate the Tangent
+    # and throws a MethodError. See Optimization.jl + AutoMooncake + ComponentArrays
+    # repro that surfaced this.
+
+    # (a) Flat-Array-backed CV: tangent is
+    #     `Tangent{@NamedTuple{data::Vector{P}, axes::NoTangent}}`.
+    let
+        cv = ComponentArray(a = randn(5), b = randn(3))
+        t = Mooncake.zero_tangent(cv)
+        copyto!(t.fields.data, 1:8)
+        out = similar(cv)
+        copyto!(out, t)
+        @test getdata(out) == collect(1.0:8.0)
+        @test out.a == [1.0, 2.0, 3.0, 4.0, 5.0]
+        @test out.b == [6.0, 7.0, 8.0]
+    end
+
+    # (b) ComponentMatrix (2D underlying storage) — same flat-Array signature, since
+    #     `Matrix{P} <: Array{P}`.
+    let
+        cm = ComponentMatrix(zeros(2, 3), Axis(r = 1:2), Axis(c = 1:3))
+        t = Mooncake.zero_tangent(cm)
+        copyto!(t.fields.data, reshape(1.0:6.0, 2, 3))
+        out = similar(cm)
+        copyto!(out, t)
+        @test getdata(out) == reshape(1.0:6.0, 2, 3)
+    end
+
+    # (c) SubArray-backed CV (from `getproperty` on a nested parent): tangent nests
+    #     a Tangent that mirrors the SubArray's `(parent, indices, offset1, stride1)`
+    #     fields. As with the symmetric `_increment_subarray_fdata!` path, copy is only
+    #     well-defined when the view fully covers its parent (because the SubArray
+    #     indices are not recoverable from Mooncake tangent shape alone).
+    let
+        # Single-top-level-component CV: the inner sub-CV's view spans the entire
+        # parent storage, so the full-cover guard succeeds.
+        parent_cv = ComponentArray(u = ComponentArray(a = randn(3), b = randn(2)))
+        sub_cv = parent_cv.u
+        @assert sub_cv isa ComponentVector{Float64, <:SubArray}
+        t = Mooncake.zero_tangent(sub_cv)
+        copyto!(t.fields.data.fields.parent, 1:5)
+        out = similar(sub_cv)              # flat-Array-backed CV
+        @assert out isa ComponentVector{Float64, <:Array}
+        copyto!(out, t)
+        @test getdata(out) == collect(1.0:5.0)
+        @test out.a == [1.0, 2.0, 3.0]
+        @test out.b == [4.0, 5.0]
+    end
+
+    # (d) Partial-cover SubArray-backed CV: the inner sub-CV views only part of the
+    #     parent storage. The full-cover guard fires with a clear ArgumentError rather
+    #     than a confusing MethodError or silently producing wrong gradients.
+    let
+        parent_cv = ComponentArray(
+            u = ComponentArray(a = randn(3), b = randn(2)), v = randn(5),
+        )
+        sub_cv = parent_cv.u
+        @assert sub_cv isa ComponentVector{Float64, <:SubArray}
+        t = Mooncake.zero_tangent(sub_cv)
+        out = similar(sub_cv)
+        @test_throws ArgumentError copyto!(out, t)
+    end
 end