Merge pull request #18 from mcabbott/master

Faster sum(BroadcastArray)
JuliaArrays · Jan 24, 2019 · 5d9754a · 5d9754a
2 parents 4559ba2 + 5c99cbb
commit 5d9754a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/lazybroadcasting.jl b/src/lazybroadcasting.jl
@@ -37,6 +37,10 @@ getindex(B::BroadcastArray{<:Any,1}, kr::AbstractVector{<:Integer}) =
 
 copy(bc::Broadcasted{<:LazyArrayStyle}) = BroadcastArray(bc)
 
+# issue 16: sum(b, dims=(1,2,3)) faster than sum(b)
+Base._sum(b::BroadcastArray{T,N}, ::Colon) where {T,N} = first(Base._sum(b, ntuple(identity, N)))
+Base._prod(b::BroadcastArray{T,N}, ::Colon) where {T,N} = first(Base._prod(b, ntuple(identity, N)))
+
 BroadcastStyle(::Type{<:BroadcastArray{<:Any,N}}) where N = LazyArrayStyle{N}()
 BroadcastStyle(L::LazyArrayStyle{N}, ::StaticArrayStyle{N}) where N = L
 BroadcastStyle(::StaticArrayStyle{N}, L::LazyArrayStyle{N})  where N = L

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -171,6 +171,9 @@ end
     @test C == A .+ 2
     D = BroadcastArray(+, A, C)
     @test D == A + C
+
+    @test sum(B) ≈ sum(exp, A)
+    @test sum(C) ≈ sum(A .+ 2)
 
     x = Vcat([3,4], [1,1,1,1,1], 1:3)
     @test x .+ (1:10) isa Vcat