From 3a26ce9be6692e88c98321fa4e1c98ea26861e4b Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Fri, 20 Dec 2019 21:54:48 -0800 Subject: [PATCH 01/11] Implement accumulate and friends --- src/mapreduce.jl | 40 ++++++++++++++++++++++++++++++++++++++++ test/accumulate.jl | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 test/accumulate.jl diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 0595a3b8..e871a8a6 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -285,3 +285,43 @@ end @inbounds return similar_type(a, T, Size($Snew))(tuple($(exprs...))) end end + +struct _InitialValue end + +_maybeval(dims::Integer) = Val(Int(dims)) +_maybeval(dims) = dims +_valof(::Val{D}) where D = D + +@inline Base.accumulate(op::F, a::StaticVector; dims = :, init = _InitialValue()) where {F} = + _accumulate(op, a, _maybeval(dims), init) + +@inline Base.accumulate(op::F, a::StaticArray; dims, init = _InitialValue()) where {F} = + _accumulate(op, a, _maybeval(dims), init) + +@inline function _accumulate(op::F, a::StaticArray, dims::Union{Val,Colon}, init) where {F} + # Adjoin the initial value to `op`: + rf(x, y) = x isa _InitialValue ? y : op(x, y) + + # StaticArrays' `reduce` is `foldl`: + results = _reduce( + a, + dims, + (init = (similar_type(a, Union{}, Size(0))(), init),), + ) do (ys, acc), x + y = rf(acc, x) + (vcat(ys, SA[y]), y) + end + dims === (:) && return first(results) + + ys = map(first, results) + data = _map(a, CartesianIndices(a)) do _, CI + D = _valof(dims) + I = Tuple(CI) + J = Base.setindex(I, 1, D) + ys[J...][I[D]] + end + return similar_type(a, eltype(data))(data) +end + +@inline Base.cumsum(a::StaticArray; kw...) = accumulate(Base.add_sum, a; kw...) +@inline Base.cumprod(a::StaticArray; kw...) = accumulate(Base.mul_prod, a; kw...) diff --git a/test/accumulate.jl b/test/accumulate.jl new file mode 100644 index 00000000..d3f8a6e8 --- /dev/null +++ b/test/accumulate.jl @@ -0,0 +1,41 @@ +using StaticArrays, Test + +@testset "accumulate" begin + @testset "cumsum(::$label)" for (label, T) in [ + # label, T + ("SVector", SVector), + ("MVector", MVector), + ("SizedVector", SizedVector{3}), + ] + a = T(SA[1, 2, 3]) + @test cumsum(a) == cumsum(collect(a)) + @test cumsum(a) isa similar_type(a) + @inferred cumsum(a) + end + + @testset "cumsum(::$label; dims=2)" for (label, T) in [ + # label, T + ("SMatrix", SMatrix), + ("MMatrix", MMatrix), + ("SizedMatrix", SizedMatrix{3,2}), + ] + a = T(SA[1 2; 3 4; 5 6]) + @test cumsum(a; dims = 2) == cumsum(collect(a); dims = 2) + @test cumsum(a; dims = 2) isa similar_type(a) + @inferred cumsum(a; dims = Val(2)) + end + + @testset "cumsum(a::SArray; dims=$i); ndims(a) = $d" for d in 1:4, i in 1:d + shape = Tuple(1:d) + a = similar_type(SArray, Int, Size(shape))(1:prod(shape)) + @test cumsum(a; dims = i) == cumsum(collect(a); dims = i) + @test cumsum(a; dims = i) isa SArray + @inferred cumsum(a; dims = Val(i)) + end + + @testset "cumprod" begin + a = SA[1, 2, 3] + @test cumprod(a)::SArray == cumprod(collect(a)) + @inferred cumprod(a) + end +end From 801d95c3715ed0e955c64e17819c700fbc74c1ae Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Sat, 21 Dec 2019 00:05:51 -0800 Subject: [PATCH 02/11] Run tests for accumulate --- test/runtests.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/runtests.jl b/test/runtests.jl index eee41e04..34b508fc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -33,6 +33,7 @@ include("abstractarray.jl") include("indexing.jl") include("initializers.jl") Random.seed!(42); include("mapreduce.jl") +Random.seed!(42); include("accumulate.jl") Random.seed!(42); include("arraymath.jl") include("broadcast.jl") include("linalg.jl") From 4985bcf090c5553c748e7e7ed0cf9c8c8491fe03 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Sat, 21 Dec 2019 00:37:57 -0800 Subject: [PATCH 03/11] Skip inference tests in Julia 1.1 --- test/accumulate.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/accumulate.jl b/test/accumulate.jl index d3f8a6e8..cdeca68a 100644 --- a/test/accumulate.jl +++ b/test/accumulate.jl @@ -22,6 +22,7 @@ using StaticArrays, Test a = T(SA[1 2; 3 4; 5 6]) @test cumsum(a; dims = 2) == cumsum(collect(a); dims = 2) @test cumsum(a; dims = 2) isa similar_type(a) + v"1.1" <= VERSION < v"1.2" && continue @inferred cumsum(a; dims = Val(2)) end @@ -30,6 +31,7 @@ using StaticArrays, Test a = similar_type(SArray, Int, Size(shape))(1:prod(shape)) @test cumsum(a; dims = i) == cumsum(collect(a); dims = i) @test cumsum(a; dims = i) isa SArray + v"1.1" <= VERSION < v"1.2" && continue @inferred cumsum(a; dims = Val(i)) end From 4ca01444f4602cc656ab467c38d7acaef15839c7 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 22:28:53 -0800 Subject: [PATCH 04/11] Update src/mapreduce.jl Co-Authored-By: Chris Foster --- src/mapreduce.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index e871a8a6..f93f0928 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -309,7 +309,7 @@ _valof(::Val{D}) where D = D (init = (similar_type(a, Union{}, Size(0))(), init),), ) do (ys, acc), x y = rf(acc, x) - (vcat(ys, SA[y]), y) + (push(ys, y), y) end dims === (:) && return first(results) From df1caa89ec95c482c32294a67f83c96ac3c173c4 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 22:06:48 -0800 Subject: [PATCH 05/11] Rename: _maybeval -> _maybe_val --- src/mapreduce.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index f93f0928..d38636ac 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -288,15 +288,15 @@ end struct _InitialValue end -_maybeval(dims::Integer) = Val(Int(dims)) -_maybeval(dims) = dims +_maybe_val(dims::Integer) = Val(Int(dims)) +_maybe_val(dims) = dims _valof(::Val{D}) where D = D @inline Base.accumulate(op::F, a::StaticVector; dims = :, init = _InitialValue()) where {F} = - _accumulate(op, a, _maybeval(dims), init) + _accumulate(op, a, _maybe_val(dims), init) @inline Base.accumulate(op::F, a::StaticArray; dims, init = _InitialValue()) where {F} = - _accumulate(op, a, _maybeval(dims), init) + _accumulate(op, a, _maybe_val(dims), init) @inline function _accumulate(op::F, a::StaticArray, dims::Union{Val,Colon}, init) where {F} # Adjoin the initial value to `op`: From df87a52eae3528d24e2865291d4455d063fa9924 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 22:28:12 -0800 Subject: [PATCH 06/11] Explain how `_map` is used from `_accumulate` --- src/mapreduce.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index d38636ac..2857a694 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -314,10 +314,13 @@ _valof(::Val{D}) where D = D dims === (:) && return first(results) ys = map(first, results) + # Now map over all indices of `a`. Since `_map` needs at least + # one `StaticArray` to be passed, we pass `a` here, even though + # the values of `a` are not used. data = _map(a, CartesianIndices(a)) do _, CI D = _valof(dims) I = Tuple(CI) - J = Base.setindex(I, 1, D) + J = setindex(I, 1, D) ys[J...][I[D]] end return similar_type(a, eltype(data))(data) From b665233c2e0a359e6d0a36b9ba24c2214b3a456f Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 22:39:25 -0800 Subject: [PATCH 07/11] Revert: (push(ys, y), y) This reverts commit 4ca01444f4602cc656ab467c38d7acaef15839c7. --- src/mapreduce.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 2857a694..b27a62ff 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -309,7 +309,7 @@ _valof(::Val{D}) where D = D (init = (similar_type(a, Union{}, Size(0))(), init),), ) do (ys, acc), x y = rf(acc, x) - (push(ys, y), y) + (vcat(ys, SA[y]), y) end dims === (:) && return first(results) From 3db20efb479420c44ffef26e4a846ad92cde6cb5 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 22:40:50 -0800 Subject: [PATCH 08/11] Comment on why we use `vcat` --- src/mapreduce.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index b27a62ff..7732b08f 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -309,6 +309,8 @@ _valof(::Val{D}) where D = D (init = (similar_type(a, Union{}, Size(0))(), init),), ) do (ys, acc), x y = rf(acc, x) + # Not using `push(ys, y)` here since we need to widen element type as + # we iterate. (vcat(ys, SA[y]), y) end dims === (:) && return first(results) From 6bf5e05b721896d484c55cdf964e99997a51e259 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 13 Feb 2020 19:58:03 -0800 Subject: [PATCH 09/11] Use inference to determine element types --- src/mapreduce.jl | 12 ++++++++++++ test/accumulate.jl | 38 +++++++++++++++++++++++++++----------- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 7732b08f..733d06e1 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -299,6 +299,18 @@ _valof(::Val{D}) where D = D _accumulate(op, a, _maybe_val(dims), init) @inline function _accumulate(op::F, a::StaticArray, dims::Union{Val,Colon}, init) where {F} + if isempty(a) + if init isa _InitialValue + # Deliberately not using `return_type` here, since this `eltype` is + # exact for the singleton element case (i.e., `op` will not be called). + return similar_type(a)() + else + # Using the type that _would_ be used if `size(a, dims) == 1`: + T = return_type(op, Tuple{typeof(init), eltype(a)}) + return similar_type(a, T)() + end + end + # Adjoin the initial value to `op`: rf(x, y) = x isa _InitialValue ? y : op(x, y) diff --git a/test/accumulate.jl b/test/accumulate.jl index cdeca68a..ca0fb156 100644 --- a/test/accumulate.jl +++ b/test/accumulate.jl @@ -5,25 +5,34 @@ using StaticArrays, Test # label, T ("SVector", SVector), ("MVector", MVector), - ("SizedVector", SizedVector{3}), + ("SizedVector", SizedVector), ] - a = T(SA[1, 2, 3]) - @test cumsum(a) == cumsum(collect(a)) - @test cumsum(a) isa similar_type(a) - @inferred cumsum(a) + @testset "$label" for (label, a) in [ + ("[1, 2, 3]", T{3}(SA[1, 2, 3])), + ("[]", T{0,Int}(())), + ] + @test cumsum(a) == cumsum(collect(a)) + @test cumsum(a) isa similar_type(a) + @inferred cumsum(a) + end end @testset "cumsum(::$label; dims=2)" for (label, T) in [ # label, T ("SMatrix", SMatrix), ("MMatrix", MMatrix), - ("SizedMatrix", SizedMatrix{3,2}), + ("SizedMatrix", SizedMatrix), ] - a = T(SA[1 2; 3 4; 5 6]) - @test cumsum(a; dims = 2) == cumsum(collect(a); dims = 2) - @test cumsum(a; dims = 2) isa similar_type(a) - v"1.1" <= VERSION < v"1.2" && continue - @inferred cumsum(a; dims = Val(2)) + @testset "$label" for (label, a) in [ + ("[1 2; 3 4; 5 6]", T{3,2}(SA[1 2; 3 4; 5 6])), + ("0 x 2 matrix", T{0,2,Float64}()), + ("2 x 0 matrix", T{2,0,Float64}()), + ] + @test cumsum(a; dims = 2) == cumsum(collect(a); dims = 2) + @test cumsum(a; dims = 2) isa similar_type(a) + v"1.1" <= VERSION < v"1.2" && continue + @inferred cumsum(a; dims = Val(2)) + end end @testset "cumsum(a::SArray; dims=$i); ndims(a) = $d" for d in 1:4, i in 1:d @@ -40,4 +49,11 @@ using StaticArrays, Test @test cumprod(a)::SArray == cumprod(collect(a)) @inferred cumprod(a) end + + @testset "empty vector with init" begin + a = SA{Int}[] + right(_, x) = x + @test accumulate(right, a; init = Val(1)) === SA{Int}[] + @inferred accumulate(right, a; init = Val(1)) + end end From 5a7963dac3bf7e1cfd8bd4551219f9e902fae7a9 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Sun, 16 Feb 2020 19:11:21 -0800 Subject: [PATCH 10/11] Use reduce_empty in cumsum/cumprod for Array-compatibility --- src/mapreduce.jl | 6 ++++-- test/accumulate.jl | 7 +++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 733d06e1..d6b23b7d 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -340,5 +340,7 @@ _valof(::Val{D}) where D = D return similar_type(a, eltype(data))(data) end -@inline Base.cumsum(a::StaticArray; kw...) = accumulate(Base.add_sum, a; kw...) -@inline Base.cumprod(a::StaticArray; kw...) = accumulate(Base.mul_prod, a; kw...) +@inline Base.cumsum(a::StaticArray; kw...) = + accumulate(Base.add_sum, a; init = Base.reduce_empty(Base.add_sum, eltype(a)), kw...) +@inline Base.cumprod(a::StaticArray; kw...) = + accumulate(Base.mul_prod, a; init = Base.reduce_empty(Base.mul_prod, eltype(a)), kw...) diff --git a/test/accumulate.jl b/test/accumulate.jl index ca0fb156..97a23e7d 100644 --- a/test/accumulate.jl +++ b/test/accumulate.jl @@ -15,6 +15,9 @@ using StaticArrays, Test @test cumsum(a) isa similar_type(a) @inferred cumsum(a) end + @test eltype(cumsum(T{0,Int8}(()))) == eltype(cumsum(Int8[])) + @test eltype(cumsum(T{1,Int8}((1)))) == eltype(cumsum(Int8[1])) + @test eltype(cumsum(T{2,Int8}((1, 2)))) == eltype(cumsum(Int8[1, 2])) end @testset "cumsum(::$label; dims=2)" for (label, T) in [ @@ -48,6 +51,10 @@ using StaticArrays, Test a = SA[1, 2, 3] @test cumprod(a)::SArray == cumprod(collect(a)) @inferred cumprod(a) + + @test eltype(cumsum(SA{Int8}[])) == eltype(cumsum(Int8[])) + @test eltype(cumsum(SA{Int8}[1])) == eltype(cumsum(Int8[1])) + @test eltype(cumsum(SA{Int8}[1, 2])) == eltype(cumsum(Int8[1, 2])) end @testset "empty vector with init" begin From 8db9ad60db0ec423bde1ed009a1d3efb318ccb21 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Sun, 16 Feb 2020 19:31:13 -0800 Subject: [PATCH 11/11] Use reduce_first instead of reduce_empty --- src/mapreduce.jl | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index d6b23b7d..b1c95636 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -299,21 +299,14 @@ _valof(::Val{D}) where D = D _accumulate(op, a, _maybe_val(dims), init) @inline function _accumulate(op::F, a::StaticArray, dims::Union{Val,Colon}, init) where {F} + # Adjoin the initial value to `op`: + rf(x, y) = x isa _InitialValue ? Base.reduce_first(op, y) : op(x, y) + if isempty(a) - if init isa _InitialValue - # Deliberately not using `return_type` here, since this `eltype` is - # exact for the singleton element case (i.e., `op` will not be called). - return similar_type(a)() - else - # Using the type that _would_ be used if `size(a, dims) == 1`: - T = return_type(op, Tuple{typeof(init), eltype(a)}) - return similar_type(a, T)() - end + T = return_type(rf, Tuple{typeof(init), eltype(a)}) + return similar_type(a, T)() end - # Adjoin the initial value to `op`: - rf(x, y) = x isa _InitialValue ? y : op(x, y) - # StaticArrays' `reduce` is `foldl`: results = _reduce( a, @@ -340,7 +333,5 @@ _valof(::Val{D}) where D = D return similar_type(a, eltype(data))(data) end -@inline Base.cumsum(a::StaticArray; kw...) = - accumulate(Base.add_sum, a; init = Base.reduce_empty(Base.add_sum, eltype(a)), kw...) -@inline Base.cumprod(a::StaticArray; kw...) = - accumulate(Base.mul_prod, a; init = Base.reduce_empty(Base.mul_prod, eltype(a)), kw...) +@inline Base.cumsum(a::StaticArray; kw...) = accumulate(Base.add_sum, a; kw...) +@inline Base.cumprod(a::StaticArray; kw...) = accumulate(Base.mul_prod, a; kw...)