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..42aa2661 100644 --- a/ext/ComponentArraysMooncakeExt.jl +++ b/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 diff --git a/test/autodiff/autodiff_tests.jl b/test/autodiff/autodiff_tests.jl index b1ef83b6..277c222a 100644 --- a/test/autodiff/autodiff_tests.jl +++ b/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