Allow init argument for sum

jump-dev · Oct 10, 2024 · a548208 · a548208
1 parent 93c7e3b
commit a548208
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/src/dispatch.jl b/src/dispatch.jl
@@ -13,8 +13,8 @@
 
 abstract type AbstractMutable end
 
-function Base.sum(a::AbstractArray{<:AbstractMutable})
-    return operate(sum, a)
+function Base.sum(a::AbstractArray{<:AbstractMutable}; kws...)
+    return operate(sum, a; kws...)
 end
 
 # When doing `x'y` where the elements of `x` and/or `y` are arrays, redirecting

diff --git a/src/reduce.jl b/src/reduce.jl
@@ -50,11 +50,10 @@ function fused_map_reduce(op::F, args::Vararg{Any,N}) where {F<:Function,N}
     return accumulator
 end
 
-function operate(::typeof(sum), a::AbstractArray)
-    return mapreduce(
-        identity,
-        add!!,
-        a;
-        init = zero(promote_operation(+, eltype(a), eltype(a))),
-    )
+function operate(
+    ::typeof(sum),
+    a::AbstractArray;
+    init = zero(promote_operation(+, eltype(a), eltype(a))),
+)
+    return mapreduce(identity, add!!, a; init)
 end