diff --git a/Project.toml b/Project.toml index 277f12aeb..faf9a412c 100644 --- a/Project.toml +++ b/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" diff --git a/src/rulesets/Base/mapreduce.jl b/src/rulesets/Base/mapreduce.jl index e8b82b32b..a9cb8013c 100644 --- a/src/rulesets/Base/mapreduce.jl +++ b/src/rulesets/Base/mapreduce.jl @@ -2,19 +2,21 @@ ##### `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) @@ -22,12 +24,13 @@ function rrule(::typeof(sum), x::Tuple) 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 @@ end ##### `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) @@ -79,6 +95,7 @@ function rrule( f::F, xs::AbstractArray{T}; dims = :, + init = _InitialValue(), ) where {F,T} project = ProjectTo(xs) @@ -91,7 +108,8 @@ function rrule( 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 @@ function rrule( # 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 @@ function frule( ::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 @@ -180,8 +204,10 @@ function rrule( ::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 @@ end # 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 diff --git a/test/rulesets/Base/mapreduce.jl b/test/rulesets/Base/mapreduce.jl index 4175fae2d..4418d15a7 100644 --- a/test/rulesets/Base/mapreduce.jl +++ b/test/rulesets/Base/mapreduce.jl @@ -15,6 +15,10 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end # 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 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end @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 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end @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 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end 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,6 +115,9 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end @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 @@ -79,6 +125,12 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end 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