sum rules do not match the signatures of the Base methods #868

The sum rules don't match the signatures of the Base methods they cover. If a call matches a rule's positional types but not its keywords, the rule throws a MethodError instead of returning nothing and leaving the call to AD. This is the same kind of bug as #768.

Base methods (same on Julia 1.10 and 1.13)

sum(a; kw...)                                 # -> mapreduce -> mapfoldl: accepts `init`
sum(a::AbstractArray; dims=:, kw...)          # accepts `dims`, `init`
sum(f, a; kw...)                              # accepts `init`
sum(f, a::AbstractArray; dims=:, kw...)       # accepts `dims`, `init`
sum!(r::AbstractArray, A::AbstractArray; init=true)
sum!(f, r::AbstractArray, A::AbstractArray; init=true)

(The specialisations for ranges, Bool, BigInt and BigFloat are faster implementations of the same function. Bool is already @non_differentiable.)

Rules that are missing a keyword Base accepts

None of the sum rules accept init:

• frule/rrule for sum(x::AbstractArray; dims=:)
• frule/rrule for sum(x::Tuple), which covers part of Base's generic sum(a; kw...) but takes no keywords
• frule for sum!(y, x). With init=false the output accumulates into y, so the tangent would be sum!(ẏ, ẋ; init=false)
• rrules for sum(f, x::Tuple) and sum(f, x::AbstractArray; dims=:)
• frule/rrule for sum(abs2, x::AbstractArray; dims=:)

Example on current main:

julia> frule((NoTangent(), ones(3)), sum, rand(3); init=0.0)
ERROR: MethodError: no method matching frule(::Tuple{NoTangent, Vector{Float64}}, ::typeof(sum), ::Vector{Float64}; init::Float64)

julia> rrule(sum, rand(3); init=0.0)
ERROR: MethodError: no method matching rrule(::typeof(sum), ::Vector{Float64}; init::Float64)

Proposal

Make every sum/sum! rule accept exactly the arguments and keywords of the Base method it covers. Treat init as non-differentiable (no tangent), as the foldl rrule already does (src/rulesets/Base/mapreduce.jl, "it does not return a gradient for init").

There can't be a rule for generic iterables: the tangent of a Generator, Dict and so on is structural rather than elementwise, so sum(ẋ) would be meaningless. Those inputs should keep falling back to AD (see #768 and #867).