Let count, sum, and prod take an init kwarg (#1281)

* let `count` take an `init` kwarg

 - `Base.count` allows the `init` kwarg since v1.6
 - bump StaticArrays to 1.9.8

* add `init` kwarg to `sum` and `prod` as well (fix #1119)

* run invalidations action on julia lts

* Move `solve.jl` tests to group B to hopefully OOM issues on Julia 1.6 CI

---------

Co-authored-by: Mateusz Baran <mateuszbaran89@gmail.com>

Author: thchr

.github/workflows/Invalidations.yml

@@ -18,7 +18,7 @@ jobs:
     steps:
     - uses: julia-actions/setup-julia@v2
       with:
-        version: '1'
+        version: 'lts'
     - uses: actions/checkout@v4
     - uses: julia-actions/julia-buildpkg@v1
     - uses: julia-actions/julia-invalidations@v1

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.9.7"
+version = "1.9.8"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/mapreduce.jl

@@ -284,16 +284,16 @@ reduce(::typeof(hcat), A::StaticArray{<:Tuple,<:StaticVecOrMatLike}) =
 # TODO: change to use Base.reduce_empty/Base.reduce_first
 @inline iszero(a::StaticArray{<:Tuple,T}) where {T} = reduce((x,y) -> x && iszero(y), a, init=true)
 
-@inline sum(a::StaticArray{<:Tuple,T}; dims=:) where {T} = _reduce(+, a, dims)
-@inline sum(f, a::StaticArray{<:Tuple,T}; dims=:) where {T} = _mapreduce(f, +, dims, _InitialValue(), Size(a), a)
-@inline sum(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:) where {T} = _mapreduce(f, +, dims, _InitialValue(), Size(a), a) # avoid ambiguity
+@inline sum(a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _reduce(+, a, dims, init)
+@inline sum(f, a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _mapreduce(f, +, dims, init, Size(a), a)
+@inline sum(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _mapreduce(f, +, dims, init, Size(a), a) # avoid ambiguity
 
-@inline prod(a::StaticArray{<:Tuple,T}; dims=:) where {T} = _reduce(*, a, dims)
-@inline prod(f, a::StaticArray{<:Tuple,T}; dims=:) where {T} = _mapreduce(f, *, dims, _InitialValue(), Size(a), a)
-@inline prod(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:) where {T} = _mapreduce(f, *, dims, _InitialValue(), Size(a), a)
+@inline prod(a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _reduce(*, a, dims, init)
+@inline prod(f, a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _mapreduce(f, *, dims, init, Size(a), a)
+@inline prod(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:, init=_InitialValue()) where {T} = _mapreduce(f, *, dims, init, Size(a), a)
 
-@inline count(a::StaticArray{<:Tuple,Bool}; dims=:) = _reduce(+, a, dims)
-@inline count(f, a::StaticArray; dims=:) = _mapreduce(x->f(x)::Bool, +, dims, _InitialValue(), Size(a), a)
+@inline count(a::StaticArray{<:Tuple,Bool}; dims=:, init=0) = _reduce(+, a, dims, init)
+@inline count(f, a::StaticArray; dims=:, init=0) = _mapreduce(x->f(x)::Bool, +, dims, init, Size(a), a)
 
 @inline all(a::StaticArray{<:Tuple,Bool}; dims=:) = _reduce(&, a, dims, true)  # non-branching versions
 @inline all(f::Function, a::StaticArray; dims=:) = _mapreduce(x->f(x)::Bool, &, dims, true, Size(a), a)

test/mapreduce.jl

@@ -130,16 +130,22 @@ using Statistics: mean
         @test sum(sa, dims=Val(2)) === RSArray2(sum(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 sum(sa, init=2) == sum(a, init=2) ≈ sum(sa) + 2 # Float64 is non-associative
+        @test sum(sb, init=2) == sum(b, init=2) == sum(sb) + 2
 
         @test prod(sa) === prod(a)
         @test prod(abs2, sa) === prod(abs2, a)
         @test prod(sa, dims=Val(2)) === RSArray2(prod(a, dims=2))
         @test prod(abs2, sa, dims=Val(2)) === RSArray2(prod(abs2, a, dims=2))
+        @test prod(sa, init=2) == prod(a, init=2) ≈ 2*prod(sa) # Float64 is non-associative
+        @test prod(sb, init=2) == prod(b, init=2) == 2*prod(sb)
 
         @test count(sb) === count(b)
         @test count(x->x>0, sa) === count(x->x>0, a)
         @test count(sb, dims=Val(2)) === RSArray2(reshape([count(b[i,:,k]) for i = 1:I, k = 1:K], (I,1,K)))
         @test count(x->x>0, sa, dims=Val(2)) === RSArray2(reshape([count(x->x>0, a[i,:,k]) for i = 1:I, k = 1:K], (I,1,K)))
+        @test count(sb, init=3) == count(b, init=3) == count(sb) + 3
+        @test count(x->x>0, sa, init=-2) == count(x->x>0, a, init=-2) == count(x->x>0, sa) - 2
 
         @test all(sb) === all(b)
         @test all(x->x>0, sa) === all(x->x>0, a)

test/runtests.jl

@@ -67,7 +67,6 @@ if TEST_GROUP ∈ ["", "all", "group-A"]
     addtests("det.jl")
     addtests("inv.jl")
     addtests("pinv.jl")
-    addtests("solve.jl")
 
     # special logic required since we need to start a new
     # Julia process for these tests
@@ -78,6 +77,7 @@ if TEST_GROUP ∈ ["", "all", "group-A"]
 end
 
 if TEST_GROUP ∈ ["", "all", "group-B"]
+    addtests("solve.jl")
     addtests("eigen.jl")
     addtests("expm.jl")
     addtests("sqrtm.jl")