Add init keyword to sum and sum! rules

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.74.12"
+version = "1.74.13"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Base/mapreduce.jl

@@ -2,32 +2,35 @@
 ##### `sum(x)`
 #####
 
-function frule((_, ẋ), ::typeof(sum), x::Tuple)
-    return sum(x), sum(ẋ)
+function frule((_, ẋ), ::typeof(sum), x::Tuple; init=_InitialValue())
+    y = init === _InitialValue() ? sum(x) : sum(x; init=init)
+    return y, sum(ẋ)
 end
 
-function frule((_, ẋ), ::typeof(sum), x; dims=:)
-    return sum(x; dims=dims), sum(ẋ; dims=dims)
+function frule((_, ẋ), ::typeof(sum), x; dims=:, init=_InitialValue())
+    y = init === _InitialValue() ? sum(x; dims=dims) : sum(x; dims=dims, init=init)
+    return y, sum(ẋ; dims=dims)
 end
 
-function frule((_, ẏ, ẋ), ::typeof(sum!), y::AbstractArray, x::AbstractArray)
-    return sum!(y, x), sum!(ẏ, ẋ)
+function frule((_, ẏ, ẋ), ::typeof(sum!), y::AbstractArray, x::AbstractArray; init=true)
+    return sum!(y, x; init=init), sum!(ẏ, ẋ; init=init)
 end
 
-function rrule(::typeof(sum), x::Tuple)
+function rrule(::typeof(sum), x::Tuple; init=_InitialValue())
     project = ProjectTo(x)
     len = Val(length(x))
     function sum_pullback(dy_raw)
         dy = unthunk(dy_raw)
         dx = dy isa AbstractZero ? dy : ntuple(Returns(dy), len)
         return (NoTangent(), project(dx))
     end
-    return sum(x), sum_pullback
+    y = init === _InitialValue() ? sum(x) : sum(x; init=init)
+    return y, sum_pullback
 end
 
-function rrule(::typeof(sum), x::AbstractArray; dims=:)
+function rrule(::typeof(sum), x::AbstractArray; dims=:, init=_InitialValue())
     project = ProjectTo(x)
-    y = sum(x; dims=dims)
+    y = init === _InitialValue() ? sum(x; dims=dims) : sum(x; dims=dims, init=init)
     function sum_pullback(dy_raw)
         dy = unthunk(dy_raw)
         x_thunk = InplaceableThunk(
@@ -62,9 +65,22 @@
 ##### `sum(f, x)`
 #####
 
-function rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(sum), f::F, xs::Tuple) where {F}
+function rrule(
+    config::RuleConfig{>:HasReverseMode},
+    ::typeof(sum),
+    f::F,
+    xs::Tuple;
+    init=_InitialValue(),
+) where {F}
     fxs, unmap = rrule(config, map, f, xs)
-    y, unsum = rrule(config, sum, fxs)
+    # Do not forward the `_InitialValue()` sentinel: the `rrule` for `sum(fxs)`
+    # may be `@non_differentiable` (e.g. `NTuple{Bool}`), which forwards all
+    # keywords to `sum`, and `sum` rejects the sentinel.
+    y, unsum = if init === _InitialValue()
+        rrule(config, sum, fxs)
+    else
+        rrule(config, sum, fxs; init=init)
+    end
     function sum_pullback_f(dy)
         _, dfxs = unsum(dy)
         _, df, dxs = unmap(dfxs)
@@ -78,7 +94,8 @@
     ::typeof(sum),
     f::F,
     xs::AbstractArray{T};
     dims = :,
+    init = _InitialValue(),
 ) where {F,T}
     project = ProjectTo(xs)
 
@@ -91,7 +108,8 @@
             end
             return (NoTangent(), NoTangent(), project(dxs))
         end
-        return sum(f, xs; dims), sum_pullback_f1
+        y = init === _InitialValue() ? sum(f, xs; dims) : sum(f, xs; dims, init=init)
+        return y, sum_pullback_f1
     end
 
     # (There is an intermediate case, where `derivatives_given_output` needs to
@@ -100,7 +118,11 @@
     # In the general case, we need to save all the pullbacks:
     # (Here `map` or `broadcast` would fail for adjoint vectors.)
     fx_and_pullbacks = [rrule_via_ad(config, f, xᵢ) for xᵢ in xs]
-    y = sum(first, fx_and_pullbacks; dims)
+    y = if init === _InitialValue()
+        sum(first, fx_and_pullbacks; dims)
+    else
+        sum(first, fx_and_pullbacks; dims, init=init)
+    end
 
     function sum_pullback_f2(dy)
         # For arrays of arrays, we ought to protect the element against broadcasting:
@@ -162,9 +184,11 @@
     ::typeof(abs2),
     x::AbstractArray{T};
     dims=:,
+    init=_InitialValue(),
 ) where {T<:Union{Real,Complex}}
     ẋ = unthunk(Δx)
-    y = sum(abs2, x; dims=dims)
+    y = init === _InitialValue() ? sum(abs2, x; dims=dims) :
+        sum(abs2, x; dims=dims, init=init)
     ∂y = if dims isa Colon
         2 * realdot(x, ẋ)
     else
@@ -176,12 +200,14 @@
 end
 
 function rrule(
     ::typeof(sum),
     ::typeof(abs2),
     x::AbstractArray{T};
     dims=:,
+    init=_InitialValue(),
 ) where {T<:Union{Real,Complex}}
-    y = sum(abs2, x; dims=dims)
+    y = init === _InitialValue() ? sum(abs2, x; dims=dims) :
+        sum(abs2, x; dims=dims, init=init)
     function sum_abs2_pullback(ȳ)
         x_thunk = InplaceableThunk(
             dx -> dx .+= 2 .* real.(ȳ) .* x,
@@ -197,9 +223,17 @@
 # and if we don't specify what do do for one that HasReverseMode then it is ambiguous
 for Config in (RuleConfig, RuleConfig{>:HasReverseMode})
     @eval function rrule(
-        ::$Config, ::typeof(sum), ::typeof(abs2), x::AbstractArray{T}; dims=:,
+        ::$Config, ::typeof(sum), ::typeof(abs2), x::AbstractArray{T};
+        dims=:, init=_InitialValue(),
     ) where {T<:Union{Real,Complex}}
-        return rrule(sum, abs2, x; dims=dims)
+        # Do not forward the `_InitialValue()` sentinel: the `rrule` this
+        # dispatches to may be `@non_differentiable` (e.g. `AbstractArray{Bool}`),
+        # which forwards all keywords to `sum`, and `sum` rejects the sentinel.
+        return if init === _InitialValue()
+            rrule(sum, abs2, x; dims=dims)
+        else
+            rrule(sum, abs2, x; dims=dims, init=init)
+        end
     end
 end
 

test/rulesets/Base/mapreduce.jl

@@ -15,6 +15,10 @@
 
         # Boolean -- via @non_differentiable
         test_rrule(sum, (true, false, true))
+
+        # `init` keyword (#868), non-differentiable like in `foldl`
+        test_frule(sum, Tuple(rand(5)); fkwargs=(; init=rand()))
+        test_rrule(sum, Tuple(rand(5)); fkwargs=(; init=rand()))
     end
     @testset "sum(x; dims=$dims)" for dims in (:, 2, (1,3))
         # Forward
@@ -25,6 +29,18 @@
         @gpu test_rrule(sum, rand(5); fkwargs=(;dims=dims))
         @gpu test_rrule(sum, rand(ComplexF64, 2,3,4); fkwargs=(;dims=dims))
 
+        # `init` keyword (#868)
+        # NOTE: `init` must match the input precision. The plain `test_*`
+        # calls run in Float64, where finite differences are exact to 1e-9;
+        # a Float32 `init` would make `sum` round its output to Float32
+        # (e.g. `sum(x; dims=2, init=0.5f0)` is `x .+ 0.5f0`), which finite
+        # differences cannot resolve to 1e-9. The `@gpu` comparisons below
+        # downcast inputs to Float32, so they use a Float32 `init` instead.
+        test_frule(sum, rand(5); fkwargs=(;dims=dims, init=rand()))
+        test_rrule(sum, rand(5); fkwargs=(;dims=dims, init=rand()))
+        @gpu frule((NoTangent(), rand(5)), sum, rand(5); dims=dims, init=rand(Float32))
+        @gpu rrule(sum, rand(5); dims=dims, init=rand(Float32))
+
         # Structured matrices
         test_rrule(sum, rand(5)'; fkwargs=(;dims=dims))
         y, back = rrule(sum, UpperTriangular(rand(5,5)); dims=dims)
@@ -54,6 +70,26 @@
     @testset "sum!(y, x)" begin
         test_frule(sum!, rand(3), rand(3, 5))
         test_frule(sum!, rand(ComplexF64, 1, 4), rand(3, 4))
+
+        # `init` keyword (#868): `init=false` accumulates into `y`
+        test_frule(sum!, rand(3), rand(3, 5); fkwargs=(; init=true))
+        test_frule(sum!, rand(3), rand(3, 5); fkwargs=(; init=false))
+    end
+
+    @testset "sum with init (#868)" begin
+        # Examples from #868: these threw `MethodError` before the fix
+        @test frule((NoTangent(), ones(3)), sum, [1.0, 2.0, 3.0]; init=0.0) == (6.0, 3.0)
+        @test rrule(sum, [1.0, 2.0, 3.0]; init=0.0)[1] == 6.0
+        @test frule((NoTangent(), ones(3)), sum, [1.0, 2.0, 3.0]; init=10.0) == (16.0, 3.0)
+        @test rrule(sum, [1.0, 2.0, 3.0]; init=10.0)[1] == 16.0
+        # `sum!` with `init=false` accumulates the tangent into `ẏ`
+        y, ẏ = [10.0, 20.0], [1.0, 2.0]
+        @test frule((NoTangent(), ẏ, ones(2, 2)), sum!, y, [1.0 2.0; 3.0 4.0]; init=false) ==
+            ([13.0, 27.0], [3.0, 4.0])
+        # Bool-valued `f` over a Tuple: `sum(fxs)` hits `@non_differentiable`,
+        # which forwards keywords to `sum` -- the `init` sentinel must not leak
+        @test rrule(CFG, sum, x -> x > 0, (1.0, 2.0, 3.0))[1] == 3
+        @test rrule(CFG, sum, x -> x > 0, (1.0, 2.0, 3.0); init=10)[1] == 13
     end
 
     @testset "sum abs2" begin
@@ -63,6 +99,13 @@
                 x = randn(T, sizes[1:N]...)
                 @gpu test_frule(sum, abs2, x; fkwargs=(;dims=dims))
                 @gpu test_rrule(sum, abs2, x; fkwargs=(;dims=dims))
+
+                # `init` keyword (#868) -- same precision split as above:
+                # plain tests in Float64, GPU comparisons in Float32
+                test_frule(sum, abs2, x; fkwargs=(;dims=dims, init=rand()))
+                test_rrule(sum, abs2, x; fkwargs=(;dims=dims, init=rand()))
+                @gpu frule((NoTangent(), NoTangent(), x), sum, abs2, x; dims=dims, init=rand(Float32))
+                @gpu rrule(sum, abs2, x; dims=dims, init=rand(Float32))
             end
 
             # Boolean -- via @non_differentiable, test that this isn't ambiguous
@@ -72,13 +115,22 @@
 
     @testset "sum(f, xs::Tuple)" begin
         test_rrule(sum, sqrt, Tuple(rand(3)), check_inferred=false)
+
+        # `init` keyword (#868)
+        test_rrule(sum, sqrt, Tuple(rand(3)); fkwargs=(; init=rand()), check_inferred=false)
     end
 
     @testset "sum(f, xs)" begin
         # This calls back into AD
         test_rrule(sum, abs, [-4.0, 2.0, 2.0])
         test_rrule(sum, log, rand(3, 4) .+ 1)
         test_rrule(sum, cbrt, randn(5))
+
+        # `init` keyword (#868): fast path (`abs`) and general path (`cbrt`)
+        test_rrule(sum, abs, [-4.0, 2.0, 2.0]; fkwargs=(; init=rand()))
+        test_rrule(sum, cbrt, randn(5); fkwargs=(; init=rand()))
+        test_rrule(sum, abs, [-2.0 4.0; 5.0 1.9]; fkwargs=(; dims=1, init=rand()))
+        test_rrule(sum, cbrt, randn(5); fkwargs=(; dims=1, init=rand()))
         test_rrule(sum, Multiplier(2.0), [2.0, 4.0, 8.0])
 
         # Complex numbers