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

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.8.23"
+version = "1.0.0-DEV"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -10,8 +10,8 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "0.10.12"
-ChainRulesTestUtils = "0.7.9"
+ChainRulesCore = "1"
+ChainRulesTestUtils = "1"
 Compat = "3.31"
 FiniteDifferences = "0.12.8"
 StaticArrays = "1.2"

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(dX -> dX .+= view(dY, ind...), @thunk(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(
-                X̄ -> mul!(X̄, Ȳ, B', true, true),
-                @thunk(Ȳ * B'),
-            ),
-            InplaceableThunk(
-                X̄ -> mul!(X̄, A', Ȳ, true, true),
-                @thunk(A' * Ȳ),
-            )
+        dA = InplaceableThunk(
+            X̄ -> mul!(X̄, Ȳ, B', true, true),
+            @thunk(project_A(Ȳ * B')),
+        )
+        dB = InplaceableThunk(
+            X̄ -> mul!(X̄, A', Ȳ, true, true),
+            @thunk(project_B(A' * Ȳ)),
         )
+        return NoTangent(), dA, dB
     end
     return A * B, times_pullback
 end
@@ -46,18 +46,20 @@ 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(
                 X̄ -> mul!(X̄, Ȳ, vec(B'), true, true),
-                @thunk(Ȳ * vec(B')),
+                @thunk(project_A(Ȳ * vec(B'))),
             ),
             InplaceableThunk(
                 X̄ -> mul!(X̄, A', Ȳ, true, true),
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
             )
         )
     end
@@ -67,14 +69,16 @@ 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(
                 X̄ -> mul!(X̄, conj(A), Ȳ, true, true),
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
             )
         )
     end
@@ -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(
                 X̄ -> mul!(X̄, conj(A), Ȳ, true, true),
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
             ),
-            @thunk(dot(Ȳ, B)'),
+            @thunk(project_A(dot(Ȳ, B)')),
         )
     end
     return A * B, times_pullback
@@ -109,27 +115,31 @@ 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(
                 dA -> mul!(dA, Ȳ, B', true, true),
-                @thunk(Ȳ * B'),
+                @thunk(project_A(Ȳ * B')),
             ),
             InplaceableThunk(
                 dB -> mul!(dB, A', Ȳ, true, true),
-                @thunk(A' * Ȳ),
+                @thunk(project_B(A' * Ȳ)),
             )
         )
         addon = if z isa Bool
             NoTangent()
         elseif z isa Number
-            @thunk(sum(Ȳ))
+            @thunk(project_z(sum(Ȳ)))
         else
             InplaceableThunk(
                 dz -> sum!(dz, Ȳ; init=false),
-                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                @thunk(project_z(sum!(similar(z, eltype(Ȳ)), Ȳ))),
             )
         end
         (NoTangent(), matmul..., addon)
@@ -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(
             dut -> dut .+= v' .* dy,
-            @thunk(v' .* dy),
+            @thunk(project_ut(v' .* dy)),
         )
         v_thunk = InplaceableThunk(
             dv -> dv .+= ut' .* dy,
-            @thunk(ut' .* dy),
+            @thunk(project_v(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,21 +179,25 @@ 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(
                 dz -> sum!(dz, Ȳ; init=false),
-                @thunk(sum!(similar(z, eltype(Ȳ)), Ȳ)),
+                @thunk(project_z(sum!(similar(z, eltype(Ȳ)), Ȳ))),
             )
         end
         (NoTangent(), proj..., addon)
@@ -202,7 +220,7 @@ function rrule(::typeof(/), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:R
         _, dBᵀ, dAᵀ = dS_pb(unthunk(dC))
 
         ∂A = last(dA_pb(unthunk(dAᵀ)))
-        ∂B = last(dA_pb(unthunk(dBᵀ)))
+        ∂B = last(dB_pb(unthunk(dBᵀ)))
 
         (NoTangent(), ∂A, ∂B)
     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)

src/rulesets/Base/fastmath_able.jl

@@ -213,9 +213,11 @@ let
         end
 
         function rrule(::typeof(*), x::Number, y::Number)
+            project_x = ProjectTo(x)
+            project_y = ProjectTo(y)
             function times_pullback(Ω̇)
                 ΔΩ = unthunk(Ω̇)
-                return (NoTangent(),  ΔΩ * y', x' * ΔΩ)
+                return (NoTangent(),  project_x(ΔΩ * y'), project_y(x' * ΔΩ))
             end
             return x * y, times_pullback
         end