Add ProjectTo for CA (#202)

* Add ProjectTo for CA

* Fix tests: gradients for Int --> Float

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

@@ -32,24 +32,22 @@ truth = ComponentArray(a = [32, 48], x = 156)
         @test zygote_full ≈ truth
     end
 
-    # Not sure why this doesn't work in v1.2, but I don't want to drop the tests for that just
-    # for this to work
-    if VERSION ≥ v"1.6"
-        @test ComponentArray(x=4,) == Zygote.gradient(ComponentArray(x=2,)) do c
-            (;c...,).x^2
-        end[1]
-    else
-        @test_skip ComponentArray(x=4,) == Zygote.gradient(ComponentArray(x=2,)) do c
-            (;c...,).x^2
-        end[1]
-    end
+    @test ComponentArray(x=4.0,) ≈ Zygote.gradient(ComponentArray(x=2,)) do c
+        (;c...,).x^2
+    end[1]
 
     # Issue #148
     ps = ComponentArray(;bias = rand(4))
     out = Zygote.gradient(x -> sum(x.^3 .+ ps.bias), Zygote.seed(rand(4),Val(12)))[1]
     @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.