Add ProjectTo for CA

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <[email protected]>"]
-version = "0.13.9"
+version = "0.13.10"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/compat/chainrulescore.jl

@@ -1,4 +1,4 @@
-function ChainRulesCore.rrule(::typeof(getproperty), x::ComponentArray, s::Union{Symbol, Val})
+function ChainRulesCore.rrule(::typeof(getproperty), x::ComponentArray, s::Union{Symbol,Val})
     function getproperty_adjoint(Δ)
         zero_x = zero(similar(x, eltype(Δ)))
         setproperty!(zero_x, s, Δ)
@@ -8,6 +8,12 @@ function ChainRulesCore.rrule(::typeof(getproperty), x::ComponentArray, s::Union
     return getproperty(x, s), getproperty_adjoint
 end
 
-ChainRulesCore.rrule(::typeof(getdata), x::ComponentArray) = getdata(x), Δ->(ChainRulesCore.NoTangent(), ComponentArray(Δ, getaxes(x)))
+ChainRulesCore.rrule(::typeof(getdata), x::ComponentArray) = getdata(x), Δ -> (ChainRulesCore.NoTangent(), ComponentArray(Δ, getaxes(x)))
 
-ChainRulesCore.rrule(::Type{ComponentArray}, data, axes) = ComponentArray(data, axes), Δ->(ChainRulesCore.NoTangent(), getdata(Δ), ChainRulesCore.NoTangent())
+ChainRulesCore.rrule(::Type{ComponentArray}, data, axes) = ComponentArray(data, axes), Δ -> (ChainRulesCore.NoTangent(), getdata(Δ), ChainRulesCore.NoTangent())
+
+function ChainRulesCore.ProjectTo(ca::ComponentArray)
+    return ChainRulesCore.ProjectTo{ComponentArray}(; project=ChainRulesCore.ProjectTo(getdata(ca)), axes=getaxes(ca))
+end
+
+(p::ChainRulesCore.ProjectTo{ComponentArray})(dx::AbstractArray) = ComponentArray(p.project(dx), p.axes)

test/autodiff_tests.jl

@@ -50,6 +50,12 @@ truth = ComponentArray(a = [32, 48], x = 156)
     @test out isa Vector{<:ForwardDiff.Dual}
 end
 
+@testset "Projection" begin
+    gs_ca = Zygote.gradient(sum, ca)[1]
+
+    @test gs_ca isa ComponentArray
+end
+
 
 # # This is commented out because the gradient operation itself is broken due to Zygote's inability
 # # to support mutation and ComponentArray's use of mutation for construction from a NamedTuple.