rrule for Base.sum should accept init keyword

[vpuri3]

The below issue contains an MWE of Zygote erroring on the init kwarg to Base.sum. Best way to fix it would be to define an rrule for sum

FluxML/Zygote.jl#1279

[vpuri3]

cc @mcabbott

[mcabbott]

There's a rule for sum,

ChainRules.jl/src/rulesets/Base/mapreduce.jl

Lines 28 to 41 in 9adf759

	function rrule(::typeof(sum), x::AbstractArray; dims=:)
	project = ProjectTo(x)
	y = sum(x; dims=dims)
	function sum_pullback(dy_raw)
	dy = unthunk(dy_raw)
	x_thunk = InplaceableThunk(
	# Protect `dy` from broadcasting, for when `x` is an array of arrays:
	dx -> dx .+= (dims isa Colon ? Ref(dy) : dy),
	@thunk project(_unsum(x, dy, dims)) # `_unsum` handles Ref internally
	)
	return (NoTangent(), x_thunk)
	end
	return y, sum_pullback
	end

, but what it doesn't do is handle init keyword, added in Julia 1.6:

julia> sum([1 2 3; 4 5 6]; init=10)
31

julia> sum([1 2 3; 4 5 6]; init=10, dims=1)
1×3 Matrix{Int64}:
 15  17  19

also:

julia> sum(abs2, [1 2 3; 4 5 6]; init=10)
101

julia> sum(abs2, [1 2 3; 4 5 6]; init=10, dims=1)
1×3 Matrix{Int64}:
 27  39  55

[mcabbott]

Xref JuliaDiff/ChainRulesCore.jl#384 --- the lowered form is this

julia> Meta.@lower sum(x; init=10)
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = Core.tuple(:init)
│   %2 = Core.apply_type(Core.NamedTuple, %1)
│   %3 = Core.tuple(10)
│   %4 = (%2)(%3)
│   %5 = Core.kwfunc(sum)
│   %6 = (%5)(%4, sum, x)
└──      return %6
))))

but defining a rule as rrule(Core.kwfunc(f), kwargs, args...) doesn't work right now.