From d555f5701991cee68cd54523bc0d47a9cca26b78 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 18 May 2026 19:37:32 -0400 Subject: [PATCH 1/2] ext/Mooncake: bridge `copyto!(::ComponentArray, ::Mooncake.Tangent)` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Mooncake.tangent_type(ComponentVector{P,Vector{P},Axes}) === Mooncake.Tangent{@NamedTuple{data::Vector{P}, axes::NoTangent}}`. That tangent is not an `AbstractArray`, so the generic `Base.copyto!(::AbstractArray, ::Any)` fallback tries to iterate it and throws `MethodError: no method matching iterate(::Mooncake.Tangent{…})`. Optimization.jl + AutoMooncake + ComponentArrays hits this path because `DifferentiationInterface.value_and_gradient!(::AutoMooncake, …)` writes the gradient into a user-supplied `grad` buffer with an unconditional `copyto!(grad, new_grad)`. The plain-`Vector` case works only because Mooncake's tangent_type for `Vector{Float64}` is `Vector{Float64}` itself. Define `Base.copyto!` for the flat-Array-backed `ComponentArray` case so the tangent's `data` field flows directly into the ComponentArray's underlying storage. SubArray-backed CVs are unaffected (the same kind of overload could be added there once a concrete failure surfaces). Co-Authored-By: Chris Rackauckas --- Project.toml | 2 +- ext/ComponentArraysMooncakeExt.jl | 17 +++++++++++++++++ test/autodiff/autodiff_tests.jl | 18 ++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index d3cc087d..14ce1df1 100644 --- a/Project.toml +++ b/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" diff --git a/ext/ComponentArraysMooncakeExt.jl b/ext/ComponentArraysMooncakeExt.jl index 7571071f..fb473c66 100644 --- a/ext/ComponentArraysMooncakeExt.jl +++ b/ext/ComponentArraysMooncakeExt.jl @@ -112,4 +112,21 @@ 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)`. For a flat-Array-backed `ComponentArray` primal, +# `Mooncake.tangent_type` is `Mooncake.Tangent{@NamedTuple{data::Vector{P}, axes::NoTangent}}`, +# 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 that by copying the tangent's `data` field directly into the ComponentArray's +# underlying storage. +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 + end diff --git a/test/autodiff/autodiff_tests.jl b/test/autodiff/autodiff_tests.jl index b1ef83b6..6dadee05 100644 --- a/test/autodiff/autodiff_tests.jl +++ b/test/autodiff/autodiff_tests.jl @@ -204,4 +204,22 @@ 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. + let + cv = ComponentArray(a = randn(5), b = randn(3)) + t = Mooncake.zero_tangent(cv) + # Populate the tangent so we can verify copyto! actually transfers the data. + 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 end From 6da5a9d77cac029530ad6376085ec03f968fa86f Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Wed, 20 May 2026 05:03:33 -0400 Subject: [PATCH 2/2] ext/Mooncake: also bridge SubArray-backed CV tangent + add ComponentMatrix test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A `ComponentVector` obtained from `getproperty(::ComponentVector, ::Symbol)` on a nested parent is `<:SubArray`-backed, so Mooncake's `tangent_type` for it nests an inner `Tangent` that mirrors the SubArray's `(parent, indices, offset1, stride1)` fields. The flat-Array `copyto!` overload doesn't dispatch on that shape, so `DI.value_and_gradient!` would still throw a `MethodError: iterate`. Add a parallel overload for the SubArray-backed CV tangent shape, symmetric to the existing `_increment_subarray_fdata!` block. As with that path, the copy is only well-defined when the view fully covers its parent (the SubArray indices are not recoverable from Mooncake tangent shape alone); the partial-cover case throws a clear `ArgumentError` rather than silently producing wrong gradients. Also add a `ComponentMatrix` copyto! regression test — `Matrix{P} <: Array{P}` already matches the flat-Array signature, this just locks in the behavior. Co-Authored-By: Chris Rackauckas --- ext/ComponentArraysMooncakeExt.jl | 52 +++++++++++++++++++++++++++---- test/autodiff/autodiff_tests.jl | 50 ++++++++++++++++++++++++++++- 2 files changed, 95 insertions(+), 7 deletions(-) diff --git a/ext/ComponentArraysMooncakeExt.jl b/ext/ComponentArraysMooncakeExt.jl index fb473c66..42aa2661 100644 --- a/ext/ComponentArraysMooncakeExt.jl +++ b/ext/ComponentArraysMooncakeExt.jl @@ -115,12 +115,13 @@ 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)`. For a flat-Array-backed `ComponentArray` primal, -# `Mooncake.tangent_type` is `Mooncake.Tangent{@NamedTuple{data::Vector{P}, axes::NoTangent}}`, -# 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 that by copying the tangent's `data` field directly into the ComponentArray's -# underlying storage. +# `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}}, @@ -129,4 +130,43 @@ function Base.copyto!( 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 diff --git a/test/autodiff/autodiff_tests.jl b/test/autodiff/autodiff_tests.jl index 6dadee05..277c222a 100644 --- a/test/autodiff/autodiff_tests.jl +++ b/test/autodiff/autodiff_tests.jl @@ -211,10 +211,12 @@ end # 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) - # Populate the tangent so we can verify copyto! actually transfers the data. copyto!(t.fields.data, 1:8) out = similar(cv) copyto!(out, t) @@ -222,4 +224,50 @@ end @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