multiple dimensions for reduce only, not mapreduce

Author: mcabbott

src/mapreduce.jl

@@ -158,11 +158,8 @@ end
     end
 end
 
-@inline function _mapreduce(f, op, D::Tuple, init, sz::Size{S}, a::StaticArray) where {S}
-    b = _mapreduce(f, op, first(D), init, sz, a)
-    return _mapreduce(f, op, Base.tail(D), init, Size(b), b)
-end
-_mapreduce(f, op, D::Tuple{}, init, sz::Size{S}, a::StaticArray) where {S} = a
+@inline _mapreduce(f, op, D::Tuple{<:Any}, init, sz::Size{S}, a::StaticArray) where {S} =
+    _mapreduce(f, op, first(D), init, sz, a)
 
 @generated function _mapfoldl(f, op, dims::Val{D}, init,
                                ::Size{S}, a::StaticArray) where {S,D}
@@ -215,6 +212,14 @@ reduce(::typeof(hcat), A::StaticArray{<:Tuple,<:StaticVecOrMatLike}) =
 @inline _reduce(op, a::StaticArray, dims, init = _InitialValue()) =
     _mapreduce(identity, op, dims, init, Size(a), a)
 
+@inline function _reduce(op, a::StaticArray, dims::Tuple, init = _InitialValue())
+    b = _reduce(op, a, first(dims), init)
+    return _reduce(op, b, Base.tail(dims))
+end
+_reduce(op, a::StaticArray, dims::Tuple{}, ::_InitialValue) = a
+_reduce(op, a::StaticArray, dims::Tuple{}, init) = op.(init, a)
+
+
 ################
 ## (map)foldl ##
 ################

test/mapreduce.jl

@@ -114,9 +114,10 @@ using Statistics: mean
         @test sum(sa, dims=2) === RSArray2(sum(a, dims=2))
         @test sum(sa, dims=(2,)) === RSArray2(sum(a, dims=2))
         @test sum(sa, dims=Val(2)) === RSArray2(sum(a, dims=2))
+        @test sum(sa, dims=(1,3)) === RSArray13(sum(a, dims=(1,3)))
         @test sum(abs2, sa; dims=2) === RSArray2(sum(abs2, a, dims=2))
+        @test sum(abs2, sa; dims=(2,)) === RSArray2(sum(abs2, a, dims=2))
         @test sum(abs2, sa; dims=Val(2)) === RSArray2(sum(abs2, a, dims=2))
-        @test_broken sum(abs2, sa; dims=(1,3)) === RSArray13(sum(abs2, a, dims=(1,3)))
 
         @test prod(sa) === prod(a)
         @test prod(abs2, sa) === prod(abs2, a)