use ProjectTo functor everywhere, and update to v1 of CRC and CRTU

Project.toml

@@ -1,16 +1,17 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.8.21"
+version = "0.8.22"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "0.10.9"
+ChainRulesCore = "0.10.11"
 ChainRulesTestUtils = "0.7.9"
 Compat = "3.31"
 FiniteDifferences = "0.12.8"
@@ -21,8 +22,7 @@ julia = "1"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesTestUtils", "FiniteDifferences", "Random", "StaticArrays", "Test"]
+test = ["ChainRulesTestUtils", "FiniteDifferences", "Random", "Test"]

src/ChainRules.jl

@@ -6,6 +6,7 @@ using Compat
 using LinearAlgebra
 using LinearAlgebra.BLAS
 using Random
+using StaticArrays
 using Statistics
 
 # Basically everything this package does is overloading these, so we make an exception
@@ -22,6 +23,10 @@ end
 # numbers that we know commute under multiplication
 const CommutativeMulNumber = Union{Real,Complex}
 
+# define ProjectTo for StaticArrays here, not in the ChainRulesCore
+(::ProjectTo{T})(dx::AbstractArray) where {T<:StaticArray} = T(dx)
+ProjectTo(x::T) where {T<:StaticArray} = ProjectTo{T}()
+
 include("rulesets/Core/core.jl")
 
 include("rulesets/Base/utils.jl")

src/rulesets/Base/array.jl

@@ -81,10 +81,11 @@ function rrule(::typeof(hcat), Xs::Union{AbstractArray, Number}...)
     Y = hcat(Xs...)  # note that Y always has 1-based indexing, even if X isa OffsetArray
     ndimsY = Val(ndims(Y))  # this avoids closing over Y, Val() is essential for type-stability
     sizes = map(size, Xs)   # this avoids closing over Xs
+    project_Xs = map(ProjectTo, Xs)
     function hcat_pullback(ȳ)
         dY = unthunk(ȳ)
         hi = Ref(0)  # Ref avoids hi::Core.Box
-        dXs = map(sizes) do sizeX
+        dXs = map(project_Xs, sizes) do project, sizeX
             ndimsX = length(sizeX)
             lo = hi[] + 1
             hi[] += get(sizeX, 2, 1)
@@ -95,14 +96,15 @@ function rrule(::typeof(hcat), Xs::Union{AbstractArray, Number}...)
                     d > ndimsX ? 1 : (:)
                 end
             end
-            if ndimsX > 0
+            dX = if ndimsX > 0
                 # Here InplaceableThunk breaks @inferred, removed for now
                 # InplaceableThunk(@thunk(dY[ind...]), dX -> dX .+= view(dY, ind...))
                 dY[ind...]
             else
                 # This is a hack to perhaps avoid GPU scalar indexing
                 sum(view(dY, ind...))
             end
+            return project(dX)
         end
         return (NoTangent(), dXs...)
     end
@@ -141,10 +143,11 @@ function rrule(::typeof(vcat), Xs::Union{AbstractArray, Number}...)
     Y = vcat(Xs...)
     ndimsY = Val(ndims(Y))
     sizes = map(size, Xs)
+    project_Xs = map(ProjectTo, Xs)
     function vcat_pullback(ȳ)
         dY = unthunk(ȳ)
         hi = Ref(0)
-        dXs = map(sizes) do sizeX
+        dXs = map(project_Xs, sizes) do project, sizeX
             ndimsX = length(sizeX)
             lo = hi[] + 1
             hi[] += get(sizeX, 1, 1)
@@ -155,12 +158,13 @@ function rrule(::typeof(vcat), Xs::Union{AbstractArray, Number}...)
                     d > ndimsX ? 1 : (:)
                 end
             end
-            if ndimsX > 0
+            dX = if ndimsX > 0
                 # InplaceableThunk(@thunk(dY[ind...]), dX -> dX .+= view(dY, ind...))
                 dY[ind...]
             else
                 sum(view(dY, ind...))
             end
+            return project(dX)
         end
         return (NoTangent(), dXs...)
     end
@@ -195,10 +199,11 @@ function rrule(::typeof(cat), Xs::Union{AbstractArray, Number}...; dims)
     cdims = dims isa Val ? Int(_val(dims)) : dims isa Integer ? Int(dims) : Tuple(dims)
     ndimsY = Val(ndims(Y))
     sizes = map(size, Xs)
+    project_Xs = map(ProjectTo, Xs)
     function cat_pullback(ȳ)
         dY = unthunk(ȳ)
         prev = fill(0, _val(ndimsY))  # note that Y always has 1-based indexing, even if X isa OffsetArray
-        dXs = map(sizes) do sizeX
+        dXs = map(project_Xs, sizes) do project, sizeX
             ndimsX = length(sizeX)
             index = ntuple(ndimsY) do d
                 if d in cdims
@@ -210,12 +215,13 @@ function rrule(::typeof(cat), Xs::Union{AbstractArray, Number}...; dims)
             for d in cdims
                 prev[d] += get(sizeX, d, 1)
             end
-            if ndimsX > 0
+            dX = if ndimsX > 0
                 # InplaceableThunk(@thunk(dY[index...]), dX -> dX .+= view(dY, index...))
                 dY[index...]
             else
                 sum(view(dY, index...))
             end
+            return project(dX)
         end
         return (NoTangent(), dXs...)
     end
@@ -231,9 +237,10 @@ function rrule(::typeof(hvcat), rows, values::Union{AbstractArray, Number}...)
     cols = size(Y,2)
     ndimsY = Val(ndims(Y))
     sizes = map(size, values)
+    project_Vs = map(ProjectTo, values)
     function hvcat_pullback(dY)
         prev = fill(0, 2)
-        dXs = map(sizes) do sizeX
+        dXs = map(project_Vs, sizes) do project, sizeX
             ndimsX = length(sizeX)
             index = ntuple(ndimsY) do d
                 if d in (1, 2)
@@ -247,7 +254,7 @@ function rrule(::typeof(hvcat), rows, values::Union{AbstractArray, Number}...)
                 prev[2] = 0
                 prev[1] += get(sizeX, 1, 1)
             end
-            dY[index...]
+            project(dY[index...])
         end
         return (NoTangent(), NoTangent(), dXs...)
     end

src/rulesets/Base/arraymath.jl

@@ -24,19 +24,19 @@ function rrule(
     A::AbstractVecOrMat{<:CommutativeMulNumber},
     B::AbstractVecOrMat{<:CommutativeMulNumber},
 )
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
     function times_pullback(ȳ)
         Ȳ = unthunk(ȳ)
-        return (
-            NoTangent(),
-            InplaceableThunk(
-                @thunk(Ȳ * B'),
-                X̄ -> mul!(X̄, Ȳ, B', true, true)
-            ),
-            InplaceableThunk(
-                @thunk(A' * Ȳ),
-                X̄ -> mul!(X̄, A', Ȳ, true, true)
-            )
+        dA = InplaceableThunk(
+            @thunk(project_A(Ȳ * B')),
+            X̄ -> mul!(X̄, Ȳ, B', true, true)
+        )
+        dB = InplaceableThunk(
+            @thunk(project_B(A' * Ȳ)),
+            X̄ -> mul!(X̄, A', Ȳ, true, true)
         )
+        return NoTangent(), dA, dB
     end
     return A * B, times_pullback
 end
@@ -46,17 +46,19 @@ function rrule(
     A::AbstractVector{<:CommutativeMulNumber},
     B::AbstractMatrix{<:CommutativeMulNumber},
 )
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
     function times_pullback(ȳ)
         Ȳ = unthunk(ȳ)
         @assert size(B, 1) === 1   # otherwise primal would have failed.
         return (
             NoTangent(),
             InplaceableThunk(
-                @thunk(Ȳ * vec(B')),
+                @thunk(project_A(Ȳ * vec(B'))),
                 X̄ -> mul!(X̄, Ȳ, vec(B'), true, true)
             ),
             InplaceableThunk(
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
                 X̄ -> mul!(X̄, A', Ȳ, true, true)
             )
         )
@@ -67,13 +69,15 @@ end
 function rrule(
    ::typeof(*), A::CommutativeMulNumber, B::AbstractArray{<:CommutativeMulNumber}
 )
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
     function times_pullback(ȳ)
         Ȳ = unthunk(ȳ)
         return (
             NoTangent(),
-            @thunk(dot(Ȳ, B)'),
+            @thunk(project_A(dot(Ȳ, B)')),
             InplaceableThunk(
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
                 X̄ -> mul!(X̄, conj(A), Ȳ, true, true)
             )
         )
@@ -84,15 +88,17 @@ end
 function rrule(
     ::typeof(*), B::AbstractArray{<:CommutativeMulNumber}, A::CommutativeMulNumber
 )
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
     function times_pullback(ȳ)
         Ȳ = unthunk(ȳ)
         return (
             NoTangent(),
             InplaceableThunk(
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
                 X̄ -> mul!(X̄, conj(A), Ȳ, true, true)
             ),
-            @thunk(dot(Ȳ, B)'),
+            @thunk(project_A(dot(Ȳ, B)')),
         )
     end
     return A * B, times_pullback
@@ -109,26 +115,30 @@ function rrule(
         B::AbstractVecOrMat{<:CommutativeMulNumber},
         z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}},
     )
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
+    project_z = ProjectTo(z)
+
     # The useful case, mul! fused with +
     function muladd_pullback_1(ȳ)
         Ȳ = unthunk(ȳ)
         matmul = (
             InplaceableThunk(
-                @thunk(Ȳ * B'),
+                @thunk(project_A(Ȳ * B')),
                 dA -> mul!(dA, Ȳ, B', true, true)
             ),
             InplaceableThunk(
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
                 dB -> mul!(dB, A', Ȳ, true, true)
             )
         )
         addon = if z isa Bool
             NoTangent()
         elseif z isa Number
-            @thunk(sum(Ȳ))
+            @thunk(project_z(sum(Ȳ)))
         else
             InplaceableThunk(
-                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                @thunk(project_z(sum!(similar(z, eltype(Ȳ)), Ȳ))),
                 dz -> sum!(dz, Ȳ; init=false)
             )
         end
@@ -143,18 +153,22 @@ function rrule(
         v::AbstractVector{<:CommutativeMulNumber},
         z::CommutativeMulNumber,
     )
+    project_ut = ProjectTo(ut)
+    project_v = ProjectTo(v)
+    project_z = ProjectTo(z)
+
     # This case is dot(u,v)+z, but would also match signature above.
     function muladd_pullback_2(ȳ)
         dy = unthunk(ȳ)
         ut_thunk = InplaceableThunk(
-            @thunk(v' .* dy),
+            @thunk(project_ut(v' .* dy)),
             dut -> dut .+= v' .* dy
         )
         v_thunk = InplaceableThunk(
-            @thunk(ut' .* dy),
+            @thunk(project_v(ut' .* dy)),
             dv -> dv .+= ut' .* dy
         )
-        (NoTangent(), ut_thunk, v_thunk, z isa Bool ? NoTangent() : dy)
+        (NoTangent(), ut_thunk, v_thunk, z isa Bool ? NoTangent() : project_z(dy))
     end
     return muladd(ut, v, z), muladd_pullback_2
 end
@@ -165,20 +179,24 @@ function rrule(
         vt::LinearAlgebra.AdjOrTransAbsVec{<:CommutativeMulNumber},
         z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}},
     )
+    project_u = ProjectTo(u)
+    project_vt = ProjectTo(vt)
+    project_z = ProjectTo(z)
+
     # Outer product, just broadcasting
     function muladd_pullback_3(ȳ)
         Ȳ = unthunk(ȳ)
         proj = (
-            @thunk(vec(sum(Ȳ .* conj.(vt), dims=2))),
-            @thunk(vec(sum(u .* conj.(Ȳ), dims=1))'),
+            @thunk(project_u(vec(sum(Ȳ .* conj.(vt), dims=2)))),
+            @thunk(project_vt(vec(sum(u .* conj.(Ȳ), dims=1))')),
         )
         addon = if z isa Bool
             NoTangent()
         elseif z isa Number
-            @thunk(sum(Ȳ))
+            @thunk(project_z(sum(Ȳ)))
         else
             InplaceableThunk(
-                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                @thunk(project_z(sum!(similar(z, eltype(Ȳ)), Ȳ))),
                 dz -> sum!(dz, Ȳ; init=false)
             )
         end
@@ -214,6 +232,9 @@ end
 #####
 
 function rrule(::typeof(\), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:Real})
+    project_A = ProjectTo(A)
+    project_B = ProjectTo(B)
+
     Y = A \ B
     function backslash_pullback(ȳ)
         Ȳ = unthunk(ȳ)
@@ -222,9 +243,9 @@ function rrule(::typeof(\), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:R
             Ā = -B̄ * Y'
             Ā = add!!(Ā, (B - A * Y) * B̄' / A')
             Ā = add!!(Ā, A' \ Y * (Ȳ' - B̄'A))
-            Ā
+            project_A(Ā)
         end
-        ∂B = @thunk A' \ Ȳ
+        ∂B = @thunk project_B(A' \ Ȳ)
         return NoTangent(), ∂A, ∂B
     end
     return Y, backslash_pullback

src/rulesets/Base/base.jl

@@ -53,9 +53,12 @@ function rrule(::Type{T}, x::Real) where {T<:Complex}
     return (T(x), Complex_pullback)
 end
 function rrule(::Type{T}, x::Number, y::Number) where {T<:Complex}
+    project_x = ProjectTo(x)
+    project_y = ProjectTo(y)
+
     function Complex_pullback(Ω̄)
         ΔΩ = unthunk(Ω̄)
-        return (NoTangent(), real(ΔΩ), imag(ΔΩ))
+        return (NoTangent(), project_x(real(ΔΩ)), project_y(imag(ΔΩ)))
     end
     return (T(x, y), Complex_pullback)
 end

src/rulesets/Base/evalpoly.jl

@@ -13,6 +13,17 @@ if VERSION ≥ v"1.4"
     end
 
     function rrule(::typeof(evalpoly), x, p)
+        y, ys = _evalpoly_intermediates(x, p)
+        project_x = ProjectTo(x)
+        project_p = p isa Tuple ? identity : ProjectTo(p)
+        function evalpoly_pullback(Δy)
+            ∂x, ∂p = _evalpoly_back(x, p, ys, Δy)
+            return NoTangent(), project_x(∂x), project_p(∂p)
+        end
+        return y, evalpoly_pullback
+    end
+
+    function rrule(::typeof(evalpoly), x, p::Vector{<:Matrix}) # does not type infer with ProjectTo
         y, ys = _evalpoly_intermediates(x, p)
         function evalpoly_pullback(Δy)
             ∂x, ∂p = _evalpoly_back(x, p, ys, Δy)