Generalize and test the rrules for sum (#217)

* Release type constraints

* Update rrule for sum abs2

* Test sum for complex

* Test sum abs2

* Generalize sum frule

* Add sum abs2 frule

* Bump version number

* Apply suggestions from code review

Co-authored-by: willtebbutt <[email protected]>

* Add fall back for pre-Julia 1.2

Co-authored-by: willtebbutt <[email protected]>

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/mapreduce.jl

@@ -2,11 +2,11 @@
 ##### `sum`
 #####
 
-function frule((_, ẋ), ::typeof(sum), x)
-    return sum(x), sum(ẋ)
+function frule((_, ẋ), ::typeof(sum), x; dims=:)
+    return sum(x; dims=dims), sum(ẋ; dims=dims)
 end
 
-function rrule(::typeof(sum), x::AbstractArray{<:Real}; dims=:)
+function rrule(::typeof(sum), x::AbstractArray{T}; dims=:) where {T<:Number}
     y = sum(sum, x; dims=dims)
     function sum_pullback(ȳ)
         # broadcasting the two works out the size no-matter `dims`
@@ -18,10 +18,33 @@ function rrule(::typeof(sum), x::AbstractArray{<:Real}; dims=:)
     return y, sum_pullback
 end
 
-function rrule(::typeof(sum), ::typeof(abs2), x::AbstractArray{<:Real}; dims=:)
+function frule(
+    (_, _, ẋ),
+    ::typeof(sum),
+    ::typeof(abs2),
+    x::AbstractArray{T};
+    dims=:,
+) where {T<:Union{Real,Complex}}
+    y = sum(abs2, x; dims=dims)
+    ∂y = if dims isa Colon
+        2 * real(dot(x, ẋ))
+    elseif VERSION ≥ v"1.2" # multi-iterator mapreduce introduced in v1.2
+        2 * mapreduce(_realconjtimes, +, x, ẋ; dims=dims)
+    else
+        2 * sum(_realconjtimes.(x, ẋ); dims=dims)
+    end
+    return y, ∂y
+end
+
+function rrule(
+    ::typeof(sum),
+    ::typeof(abs2),
+    x::AbstractArray{T};
+    dims=:,
+) where {T<:Union{Real,Complex}}
     y = sum(abs2, x; dims=dims)
     function sum_abs2_pullback(ȳ)
-        return (NO_FIELDS, DoesNotExist(), @thunk(2ȳ .* x))
+        return (NO_FIELDS, DoesNotExist(), @thunk(2 .* real.(ȳ) .* x))
     end
     return y, sum_abs2_pullback
 end

test/rulesets/Base/mapreduce.jl

@@ -1,28 +1,41 @@
 @testset "Maps and Reductions" begin
     @testset "sum" begin
-        @testset "Vector" begin
-            M = 3
-            frule_test(sum, (randn(M), randn(M)))
-            rrule_test(sum, randn(), (randn(M), randn(M)))
-        end
-        @testset "Matrix" begin
-            M, N = 3, 4
-            frule_test(sum, (randn(M, N), randn(M, N)))
-            rrule_test(sum, randn(), (randn(M, N), randn(M, N)))
-        end
-        @testset "Array{T, 3}" begin
-            M, N, P = 3, 7, 11
-            frule_test(sum, (randn(M, N, P), randn(M, N, P)))
-            rrule_test(sum, randn(), (randn(M, N, P), randn(M, N, P)))
-        end
-        @testset "keyword arguments" begin
-            n = 4
-            X = randn(n, n+1)
-            y, pullback = rrule(sum, X; dims=2)
-            ȳ = randn(size(y))
-            _, x̄_ad = pullback(ȳ)
-            x̄_fd = only(j′vp(central_fdm(5, 1), x->sum(x, dims=2), ȳ, X))
-            @test x̄_ad ≈ x̄_fd atol=1e-9 rtol=1e-9
+        sizes = (3, 4, 7)
+        @testset "dims = $dims" for dims in (:, 1)
+            fkwargs = (dims=dims,)
+            @testset "Array{$N, $T}" for N in eachindex(sizes), T in (Float64, ComplexF64)
+                s = sizes[1:N]
+                x = randn(T, s...)
+                ẋ = randn(T, s...)
+                x̄ = randn(T, s...)
+                y = sum(x; dims=dims)
+                Δy = randn(eltype(y), size(y)...)
+                frule_test(sum, (x, ẋ); fkwargs=fkwargs)
+                rrule_test(sum, Δy, (x, x̄); fkwargs=fkwargs)
+            end
         end
     end  # sum
+
+    @testset "sum abs2" begin
+        sizes = (3, 4, 7)
+        @testset "dims = $dims" for dims in (:, 1)
+            fkwargs = (dims=dims,)
+            @testset "Array{$N, $T}" for N in eachindex(sizes), T in (Float64, ComplexF64)
+                s = sizes[1:N]
+                x, ẋ, x̄ = randn(T, s...), randn(T, s...), randn(T, s...)
+                y = sum(abs2, x; dims=dims)
+                Δy = randn(eltype(y), size(y)...)
+                @testset "frule" begin
+                    # can't use frule_test here because it doesn't yet ignore nothing tangents
+                    y_ad, ẏ_ad = frule((Zero(), Zero(), ẋ), sum, abs2, x; dims=dims)
+                    @test y_ad == y
+                    ẏ_fd = jvp(_fdm, z -> sum(abs2, z; dims=dims), (x, ẋ))
+                    @test ẏ_ad ≈ ẏ_fd
+                end
+                @testset "rrule" begin
+                    rrule_test(sum, Δy, (abs2, nothing), (x, x̄); fkwargs=fkwargs)
+                end
+            end
+        end
+    end  # sum abs2
 end