Add mapreduce_layout hook for reductions on LayoutArrays

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
-version = "1.12.4"
+version = "1.13.0"
 authors = ["Sheehan Olver <solver@mac.com>"]
 
 [deps]

src/ArrayLayouts.jl

@@ -305,6 +305,10 @@ _norm(_, A, p) = invoke(norm, Tuple{Any,Real}, A, p)
 LinearAlgebra.norm(A::LayoutArray, p::Real=2) = _norm(MemoryLayout(A), A, p)
 LinearAlgebra.norm(A::SubArray{<:Any,N,<:LayoutArray}, p::Real=2) where N = _norm(MemoryLayout(A), A, p)
 
+mapreduce_layout(::Any, f, op, A, dims; kw...) = invoke(mapreduce, Tuple{Any,Any,AbstractArray}, f, op, A; dims, kw...)
+Base.mapreduce(f, op, A::LayoutArray; dims=:, kw...) = mapreduce_layout(MemoryLayout(A), f, op, A, dims; kw...)
+Base.mapreduce(f, op, A::SubArray{<:Any,N,<:LayoutArray}; dims=:, kw...) where N = mapreduce_layout(MemoryLayout(A), f, op, A, dims; kw...)
+
 
 _fill_lmul!(β, A::AbstractArray) = iszero(β) ? zero!(A) : lmul!(β, A)
 

test/test_layoutarray.jl

@@ -82,6 +82,9 @@ BroadcastStyle(::Type{<:MyVector2{T}}) where {T} = MyBroadcastStyle{1}()
         @test copy(v) == sub_materialize(v) == a[1:3]
         @test dot(v,a) == dot(v,a.A) == dot(a,v) == dot(a.A,v) == dot(v,v) == 14
         @test norm(v) == norm(a) == norm([1,2,3])
+        @test sum(v) == sum(a) == 6
+        @test sum(a; init=1) == 7
+        @test sum(a; dims=1) == [6]
 
         V = view(a',:,1:3)
         @test copy(V) == sub_materialize(V) == (a')[:,1:3]
@@ -697,6 +700,32 @@ triangulardata(A::MyUpperTriangular) = triangulardata(A.A)
     @test U / MyMatrix(A) ≈ U / A
 end
 
+struct MyReduceLayout <: MemoryLayout end
+struct MyReduceVector <: LayoutVector{Float64} end
+Base.size(::MyReduceVector) = (3,)
+Base.getindex(::MyReduceVector, k::Int) = Float64(k)
+MemoryLayout(::Type{MyReduceVector}) = MyReduceLayout()
+ArrayLayouts.mapreduce_layout(::MyReduceLayout, f, op, A, ::Colon; kw...) = -1.0
+
+@testset "mapreduce" begin
+    A = randn(5,4)
+    M = MyMatrix(A)
+    @test sum(M) ≈ sum(A)
+    @test sum(M; dims=1) ≈ sum(A; dims=1)
+    @test sum(M; dims=2) ≈ sum(A; dims=2)
+    @test sum(abs2, M) ≈ sum(abs2, A)
+    @test sum(M; init=1.0) ≈ sum(A; init=1.0)
+    @test maximum(M) == maximum(A)
+    @test extrema(M) == extrema(A)
+    @test prod(M) ≈ prod(A)
+    @test mapreduce(abs, max, M; dims=1) == mapreduce(abs, max, A; dims=1)
+    @test sum(view(M, 1:3, 2:4)) ≈ sum(A[1:3, 2:4])
+
+    v = MyReduceVector()
+    @test sum(v) == maximum(v) == reduce(+, v) == -1.0
+    @test sum(v; dims=1) == [6.0]
+end
+
 @testset "* for infinite layouts" begin
     tup = InfSymTridiagonal(), InfTridiagonal(), InfBidiagonal('U'),
         InfBidiagonal('L'),