From 6221324bbf2fd07034b1918c73e1a1a6d0a023df Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Tue, 27 Aug 2024 18:34:19 -0400 Subject: [PATCH 1/6] WIP: implement reductions as reductions Depends upon (and is a requirement for) Julia issue 55318 --- src/Statistics.jl | 158 +++++++++++++++++++++++----------------------- test/runtests.jl | 40 ++++++++++-- 2 files changed, 113 insertions(+), 85 deletions(-) diff --git a/src/Statistics.jl b/src/Statistics.jl index 6633edf0..d776b9e2 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -43,6 +43,18 @@ julia> mean(skipmissing([1, missing, 3])) """ mean(itr) = mean(identity, itr) +struct Counter{F} <: Function + f::F + n::Base.RefValue{Int} +end +Counter(f::F) where {F} = Counter{F}(f, Ref(0)) +(f::Counter)(x) = (f.n[] += 1; f.f(x)) + +struct DivOne{F} <: Function + f::F +end +(f::DivOne)(x) = f.f(x)/1 + """ mean(f, itr) @@ -59,23 +71,13 @@ julia> mean([√1, √2, √3]) ``` """ function mean(f, itr) - y = iterate(itr) - if y === nothing - return Base.mapreduce_empty_iter(f, +, itr, - Base.IteratorEltype(itr)) / 0 - end - count = 1 - value, state = y - f_value = f(value)/1 - total = Base.reduce_first(+, f_value) - y = iterate(itr, state) - while y !== nothing - value, state = y - total += _mean_promote(total, f(value)) - count += 1 - y = iterate(itr, state) + if Base.IteratorSize(itr) === Base.SizeUnknown() + g = Counter(DivOne(f)) + result = mapfoldl(g, add_mean, itr) + return result/g.n[] + else + return mapfoldl(DivOne(f), add_mean, itr)/length(itr) end - return total/count end """ @@ -180,20 +182,24 @@ mean(A::AbstractArray; dims=:) = _mean(identity, A, dims) _mean_promote(x::T, y::S) where {T,S} = convert(promote_type(T, S), y) +add_mean(x, y) = Base.add_sum(x, _mean_promote(x, y)) + +Base.reduce_empty(::typeof(add_mean), T) = Base.reduce_empty(Base.add_sum, T) +Base.mapreduce_empty(g::DivOne, ::typeof(add_mean), T) = Base.mapreduce_empty(g.f, Base.add_sum, T)/1 +Base.mapreduce_empty(g::Counter{<:DivOne}, ::typeof(add_mean), T) = Base.mapreduce_empty(g.f.f, Base.add_sum, T)/1 + + # ::Dims is there to force specializing on Colon (as it is a Function) function _mean(f, A::AbstractArray, dims::Dims=:) where Dims - isempty(A) && return sum(f, A, dims=dims)/0 if dims === (:) + result = mapreduce(DivOne(f), add_mean, A, dims=dims) n = length(A) - else - n = mapreduce(i -> size(A, i), *, unique(dims); init=1) - end - x1 = f(first(A)) / 1 - result = sum(x -> _mean_promote(x1, f(x)), A, dims=dims) - if dims === (:) return result / n else - return result ./= n + result = mapreduce(DivOne(f), add_mean, A, dims=dims) + n = prod(i -> size(A, i), unique(dims); init=1) + result ./= n + return result end end @@ -211,6 +217,7 @@ realXcY(x::Complex, y::Complex) = real(x)*real(y) + imag(x)*imag(y) var(iterable; corrected::Bool=true, mean=nothing) = _var(iterable, corrected, mean) function _var(iterable, corrected::Bool, mean) + ismissing(mean) && return missing y = iterate(iterable) if y === nothing T = eltype(iterable) @@ -252,61 +259,36 @@ function _var(iterable, corrected::Bool, mean) end end -centralizedabs2fun(m) = x -> abs2.(x - m) -centralize_sumabs2(A::AbstractArray, m) = - mapreduce(centralizedabs2fun(m), +, A) -centralize_sumabs2(A::AbstractArray, m, ifirst::Int, ilast::Int) = - Base.mapreduce_impl(centralizedabs2fun(m), +, A, ifirst, ilast) - -function centralize_sumabs2!(R::AbstractArray{S}, A::AbstractArray, means::AbstractArray) where S - # following the implementation of _mapreducedim! at base/reducedim.jl - lsiz = Base.check_reducedims(R,A) - for i in 1:max(ndims(R), ndims(means)) - if axes(means, i) != axes(R, i) - throw(DimensionMismatch("dimension $i of `mean` should have indices $(axes(R, i)), but got $(axes(means, i))")) - end - end - isempty(R) || fill!(R, zero(S)) - isempty(A) && return R - - if Base.has_fast_linear_indexing(A) && lsiz > 16 && !has_offset_axes(R, means) - nslices = div(length(A), lsiz) - ibase = first(LinearIndices(A))-1 - for i = 1:nslices - @inbounds R[i] = centralize_sumabs2(A, means[i], ibase+1, ibase+lsiz) - ibase += lsiz - end - return R - end - indsAt, indsRt = Base.safe_tail(axes(A)), Base.safe_tail(axes(R)) # handle d=1 manually - keep, Idefault = Broadcast.shapeindexer(indsRt) - if Base.reducedim1(R, A) - i1 = first(Base.axes1(R)) - @inbounds for IA in CartesianIndices(indsAt) - IR = Broadcast.newindex(IA, keep, Idefault) - r = R[i1,IR] - m = means[i1,IR] - @simd for i in axes(A, 1) - r += abs2(A[i,IA] - m) - end - R[i1,IR] = r - end - else - @inbounds for IA in CartesianIndices(indsAt) - IR = Broadcast.newindex(IA, keep, Idefault) - @simd for i in axes(A, 1) - R[i,IR] += abs2(A[i,IA] - means[i,IR]) - end - end - end - return R +struct CentralizedAbs2Fun{T,S} <: Function + mean::S end +CentralizedAbs2Fun{T}(means) where {T} = CentralizedAbs2Fun{T,typeof(means)}(means) +CentralizedAbs2Fun(means) = CentralizedAbs2Fun{typeof(means)}(means) +CentralizedAbs2Fun(means, extrude) = CentralizedAbs2Fun{eltype(means)}(Broadcast.extrude(means)) +# Division is generally costly, but Julia is typically able to constant propagate a /1 +# and simply ensure we get the type right at no cost, allowing the division in-place later +(f::CentralizedAbs2Fun)(x) = abs2.(x - f.mean)/1 +(f::CentralizedAbs2Fun{<:Any,<:Broadcast.Extruded})((i, x),) = abs2.(x - Broadcast._broadcast_getindex(f.mean, i))/1 +_doubled(x) = x+x +Base.mapreduce_empty(::CentralizedAbs2Fun{T,<:Broadcast.Extruded}, ::typeof(Base.add_sum), ::Type{Tuple{_Any,S}}) where {T<:Number, S<:Number, _Any} = _doubled(abs2(zero(T)-zero(S)))/1 +Base.mapreduce_empty(::CentralizedAbs2Fun{T,<:Broadcast.Extruded}, ::typeof(Base.add_sum), ::Type{Tuple{_Any, Union{Missing, S}}}) where {T<:Number, S<:Number, _Any} = _doubled(abs2(zero(T)-zero(S)))/1 +Base.mapreduce_empty(::CentralizedAbs2Fun{T}, ::typeof(Base.add_sum), ::Type{S}) where {T<:Number, S<:Number} = _doubled(abs2(zero(T)-zero(S)))/1 +Base.mapreduce_empty(::CentralizedAbs2Fun{T}, ::typeof(Base.add_sum), ::Type{Union{Missing, S}}) where {T<:Number, S<:Number} = _doubled(abs2(zero(T)-zero(S)))/1 + +centralize_sumabs2(A::AbstractArray, m) = + sum(CentralizedAbs2Fun(m), A) +centralize_sumabs2(A::AbstractArray, m::AbstractArray, region) = + sum(CentralizedAbs2Fun(m, true), Base.PairsArray(A), dims=region) +centralize_sumabs2!(R::AbstractArray, A::AbstractArray, means::AbstractArray) = + sum!(CentralizedAbs2Fun(means, true), R, Base.PairsArray(A)) + function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S - if isempty(A) + _checkm(R, m, ntuple(identity, Val(max(ndims(R), ndims(m))))) + if isempty(A) || length(A) == 1 && corrected fill!(R, convert(S, NaN)) else - rn = div(length(A), length(R)) - Int(corrected) + rn = prod(ntuple(d->size(R, d) == 1 ? size(A, d) : 1, Val(max(ndims(A), ndims(R))))) - Int(corrected) centralize_sumabs2!(R, A, m) R .= R .* (1 // rn) end @@ -339,15 +321,33 @@ over dimensions. In that case, `mean` must be an array with the same shape as """ varm(A::AbstractArray, m::AbstractArray; corrected::Bool=true, dims=:) = _varm(A, m, corrected, dims) -_varm(A::AbstractArray{T}, m, corrected::Bool, region) where {T} = - varm!(Base.reducedim_init(t -> abs2(t)/2, +, A, region), A, m; corrected=corrected) +_throw_mean_mismatch(A, m, region) = throw(DimensionMismatch("axes of means ($(axes(m))) does not match reduction over $(region) of $(axes(A))")) +function _checkm(A::AbstractArray, m::AbstractArray, region) + for d in 1:max(ndims(A), ndims(m)) + if d in region + size(m, d) == 1 || _throw_mean_mismatch(A, m, region) + else + axes(m, d) == axes(A, d) || _throw_mean_mismatch(A, m, region) + end + end +end +function _varm(A::AbstractArray, m, corrected::Bool, region) + _checkm(A, m, region) + rn = prod(ntuple(d->d in region ? size(A, d) : 1, Val(ndims(A)))) - Int(corrected) + R = centralize_sumabs2(A, m, region) + if rn <= 0 + R .= R ./ 0 + else + R .= R .* 1//rn # why use Rational? + end + return R +end varm(A::AbstractArray, m; corrected::Bool=true) = _varm(A, m, corrected, :) function _varm(A::AbstractArray{T}, m, corrected::Bool, ::Colon) where T - n = length(A) - n == 0 && return oftype((abs2(zero(T)) + abs2(zero(T)))/2, NaN) - return centralize_sumabs2(A, m) / (n - Int(corrected)) + rn = max(length(A) - Int(corrected), 0) + centralize_sumabs2(A, m)/rn end diff --git a/test/runtests.jl b/test/runtests.jl index d2f25412..79b4a551 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -522,9 +522,15 @@ Y = [6.0 2.0; @testset "cov with missing" begin @test cov([missing]) === cov([1, missing]) === missing @test cov([1, missing], [2, 3]) === cov([1, 3], [2, missing]) === missing - @test_throws Exception cov([1 missing; 2 3]) - @test_throws Exception cov([1 missing; 2 3], [1, 2]) - @test_throws Exception cov([1, 2], [1 missing; 2 3]) + if isdefined(Base, :_reducedim_init) + @test_broken isequal(coalesce.(cov([1 missing; 2 3]), NaN), cov([1 NaN; 2 3])) + @test_broken isequal(coalesce.(cov([1 missing; 2 3], [1, 2]), NaN), cov([1 NaN; 2 3], [1, 2])) + @test_broken isequal(coalesce.(cov([1, 2], [1 missing; 2 3]), NaN), cov([1, 2], [1 NaN; 2 3])) + else + @test isequal(coalesce.(cov([1 missing; 2 3]), NaN), cov([1 NaN; 2 3])) + @test isequal(coalesce.(cov([1 missing; 2 3], [1, 2]), NaN), cov([1 NaN; 2 3], [1, 2])) + @test isequal(coalesce.(cov([1, 2], [1 missing; 2 3]), NaN), cov([1, 2], [1 NaN; 2 3])) + end @test isequal(cov([1 2; 2 3], [1, missing]), [missing missing]') @test isequal(cov([1, missing], [1 2; 2 3]), [missing missing]) end @@ -633,9 +639,15 @@ end @test cor([missing]) === missing @test cor([1, missing]) == 1 @test cor([1, missing], [2, 3]) === cor([1, 3], [2, missing]) === missing - @test_throws Exception cor([1 missing; 2 3]) - @test_throws Exception cor([1 missing; 2 3], [1, 2]) - @test_throws Exception cor([1, 2], [1 missing; 2 3]) + if isdefined(Base, :_reducedim_init) + @test_broken isequal(coalesce.(cor([1 missing; 2 3]), NaN), cor([1 NaN; 2 3])) + @test_broken isequal(coalesce.(cor([1 missing; 2 3], [1, 2]), NaN), cor([1 NaN; 2 3], [1, 2])) + @test_broken isequal(coalesce.(cor([1, 2], [1 missing; 2 3]), NaN), cor([1, 2], [1 NaN; 2 3])) + else + @test isequal(coalesce.(cor([1 missing; 2 3]), NaN), cor([1 NaN; 2 3])) + @test isequal(coalesce.(cor([1 missing; 2 3], [1, 2]), NaN), cor([1 NaN; 2 3], [1, 2])) + @test isequal(coalesce.(cor([1, 2], [1 missing; 2 3]), NaN), cor([1, 2], [1 NaN; 2 3])) + end @test isequal(cor([1 2; 2 3], [1, missing]), [missing missing]') @test isequal(cor([1, missing], [1 2; 2 3]), [missing missing]) end @@ -1072,3 +1084,19 @@ end @test isequal(cov(Int[], my), fill(-0.0, 1, 3)) @test isequal(cor(Int[], my), fill(NaN, 1, 3)) end + +@testset "mean, var, std type stability with Missings; Issue #160" begin + @test (@inferred Missing mean(view([1, 2, missing], 1:2))) == (@inferred mean([1,2])) + @test (@inferred Missing var(view([1, 2, missing], 1:2))) == (@inferred var([1,2])) + @test (@inferred Missing std(view([1, 2, missing], 1:2))) == (@inferred std([1,2])) +end + +@testset "inexact errors; Issues #7 and #126" begin + a = [missing missing; 0 1] + @test isequal(mean(a;dims=2), [missing; 0.5;;]) + + x = [(i==3 && j==3) ? missing : i*j for i in 1:3, j in 1:4] + @test ismissing(@inferred Float64 mean(x)) + @test isequal(mean(x; dims=1), [2. 4. missing 8.]) + @test isequal(mean(x; dims=2), [2.5; 5.0; missing;;]) +end From c674ab823526d05e7a85b468065089bf563823aa Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Wed, 14 May 2025 17:50:00 -0400 Subject: [PATCH 2/6] skip the mean refactor for now --- src/Statistics.jl | 54 +++++++++++++++++++++-------------------------- test/runtests.jl | 2 +- 2 files changed, 25 insertions(+), 31 deletions(-) diff --git a/src/Statistics.jl b/src/Statistics.jl index d776b9e2..b34190d5 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -43,18 +43,6 @@ julia> mean(skipmissing([1, missing, 3])) """ mean(itr) = mean(identity, itr) -struct Counter{F} <: Function - f::F - n::Base.RefValue{Int} -end -Counter(f::F) where {F} = Counter{F}(f, Ref(0)) -(f::Counter)(x) = (f.n[] += 1; f.f(x)) - -struct DivOne{F} <: Function - f::F -end -(f::DivOne)(x) = f.f(x)/1 - """ mean(f, itr) @@ -71,13 +59,23 @@ julia> mean([√1, √2, √3]) ``` """ function mean(f, itr) - if Base.IteratorSize(itr) === Base.SizeUnknown() - g = Counter(DivOne(f)) - result = mapfoldl(g, add_mean, itr) - return result/g.n[] - else - return mapfoldl(DivOne(f), add_mean, itr)/length(itr) + y = iterate(itr) + if y === nothing + return Base.mapreduce_empty_iter(f, +, itr, + Base.IteratorEltype(itr)) / 0 end + count = 1 + value, state = y + f_value = f(value)/1 + total = Base.reduce_first(+, f_value) + y = iterate(itr, state) + while y !== nothing + value, state = y + total += _mean_promote(total, f(value)) + count += 1 + y = iterate(itr, state) + end + return total/count end """ @@ -182,24 +180,20 @@ mean(A::AbstractArray; dims=:) = _mean(identity, A, dims) _mean_promote(x::T, y::S) where {T,S} = convert(promote_type(T, S), y) -add_mean(x, y) = Base.add_sum(x, _mean_promote(x, y)) - -Base.reduce_empty(::typeof(add_mean), T) = Base.reduce_empty(Base.add_sum, T) -Base.mapreduce_empty(g::DivOne, ::typeof(add_mean), T) = Base.mapreduce_empty(g.f, Base.add_sum, T)/1 -Base.mapreduce_empty(g::Counter{<:DivOne}, ::typeof(add_mean), T) = Base.mapreduce_empty(g.f.f, Base.add_sum, T)/1 - - # ::Dims is there to force specializing on Colon (as it is a Function) function _mean(f, A::AbstractArray, dims::Dims=:) where Dims + isempty(A) && return sum(f, A, dims=dims)/0 if dims === (:) - result = mapreduce(DivOne(f), add_mean, A, dims=dims) n = length(A) + else + n = mapreduce(i -> size(A, i), *, unique(dims); init=1) + end + x1 = f(first(A)) / 1 + result = sum(x -> _mean_promote(x1, f(x)), A, dims=dims) + if dims === (:) return result / n else - result = mapreduce(DivOne(f), add_mean, A, dims=dims) - n = prod(i -> size(A, i), unique(dims); init=1) - result ./= n - return result + return result ./= n end end diff --git a/test/runtests.jl b/test/runtests.jl index 79b4a551..6e9d095f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1093,7 +1093,7 @@ end @testset "inexact errors; Issues #7 and #126" begin a = [missing missing; 0 1] - @test isequal(mean(a;dims=2), [missing; 0.5;;]) + @test_broken isequal(mean(a;dims=2), [missing; 0.5;;]) x = [(i==3 && j==3) ? missing : i*j for i in 1:3, j in 1:4] @test ismissing(@inferred Float64 mean(x)) From 257e8271adcd7c08d50ce370c1bdab71485187e0 Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Thu, 30 Jul 2026 23:26:55 -0400 Subject: [PATCH 3/6] varm!: fix the shape check for array means and single-element behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-place `varm!` validated its `means` argument by passing every dimension as the reduction region to `_checkm`, which demands a singleton size along each axis. Any correctly-shaped means array — one matching the output `R` — was therefore rejected with a `DimensionMismatch`, and no test exercised `varm!` with array means. Restore the intended check: `means` must have the same axes as `R`. Also drop the special case that filled `R` with NaN whenever `A` had a single element and `corrected=true`. The out-of-place dims path yields `Inf` in that situation when the provided mean differs from the lone element. Do what `_varm` does instead: when the correction leaves no degrees of freedom, divide by zero, so a matching mean gives 0/0 = NaN and a differing mean gives Inf. Assisted-By: Claude Fable 5 --- src/Statistics.jl | 15 ++++++++++++--- test/runtests.jl | 20 ++++++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/src/Statistics.jl b/src/Statistics.jl index 9c5d0c71..766618c2 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -283,13 +283,22 @@ centralize_sumabs2!(R::AbstractArray, A::AbstractArray, means::AbstractArray) = function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S - _checkm(R, m, ntuple(identity, Val(max(ndims(R), ndims(m))))) - if isempty(A) || length(A) == 1 && corrected + for d in 1:max(ndims(R), ndims(m)) + axes(m, d) == axes(R, d) || throw(DimensionMismatch( + "dimension $d of `mean` should have indices $(axes(R, d)), but got $(axes(m, d))")) + end + if isempty(A) fill!(R, convert(S, NaN)) else rn = prod(ntuple(d->size(R, d) == 1 ? size(A, d) : 1, Val(max(ndims(A), ndims(R))))) - Int(corrected) centralize_sumabs2!(R, A, m) - R .= R .* (1 // rn) + if rn <= 0 + # Like the out-of-place `_varm` path: a corrected variance over a + # single element is 0/0 = NaN (or Inf if `m` is not that element) + R .= R ./ 0 + else + R .= R .* (1 // rn) + end end return R end diff --git a/test/runtests.jl b/test/runtests.jl index 83369667..16761dba 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -383,6 +383,26 @@ end @test isequal(var(A, dims=3), fill(NaN, 0, 1)) end + @testset "varm! with array means" begin + A = [1.0 2.0; 3.0 4.0] + @test Statistics.varm!(zeros(1, 2), A, mean(A, dims=1)) ≈ var(A, dims=1) + @test Statistics.varm!(zeros(2, 1), A, mean(A, dims=2)) ≈ var(A, dims=2) + @test Statistics.varm!(zeros(2), A, vec(mean(A, dims=2))) ≈ vec(var(A, dims=2)) + @test Statistics.varm!(zeros(1, 2), A, mean(A, dims=1); corrected=false) ≈ + var(A, dims=1, corrected=false) + # means must have the same shape as the result + @test_throws DimensionMismatch Statistics.varm!(zeros(1, 2), A, mean(A, dims=2)) + @test_throws DimensionMismatch Statistics.varm!(zeros(2, 1), A, mean(A, dims=1)) + @test_throws DimensionMismatch Statistics.varm!(zeros(2, 1), A, [1.0 2.0 3.0]) + # single-element slices match the out-of-place dims path: 0/0 = NaN + # for a matching mean and Inf otherwise (NaN sign/payload is + # platform-dependent, so compare with isequal, not ===) + B = [2.0; 4.0;;] + @test isequal(Statistics.varm!(zeros(2, 1), B, [2.0; 5.0;;]), + varm(B, [2.0; 5.0;;]; dims=2)) + @test isequal(Statistics.varm!(zeros(2, 1), B, [2.0; 5.0;;]), [NaN; Inf;;]) + end + # issue #6672 @test std(AbstractFloat[1,2,3], dims=1) == [1.0] From 9a80ac7acedc8717d7573aae40438c8e5a2cd294 Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Thu, 30 Jul 2026 23:28:28 -0400 Subject: [PATCH 4/6] variance: index the extruded means directly with inbounds The dims-variance kernel looked up each slice's mean through `Broadcast._broadcast_getindex`, which bounds-checks the means array for every element of `A`; that check accounted for most of the slowdown against the old hand-rolled loops. Compute the broadcast index explicitly and assert the access in-bounds. Assisted-By: Claude Fable 5 --- src/Statistics.jl | 20 +++++++++++++++----- test/runtests.jl | 15 +++++++++++++++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/Statistics.jl b/src/Statistics.jl index 766618c2..3686af05 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -263,11 +263,19 @@ struct CentralizedAbs2Fun{T,S} <: Function end CentralizedAbs2Fun{T}(means) where {T} = CentralizedAbs2Fun{T,typeof(means)}(means) CentralizedAbs2Fun(means) = CentralizedAbs2Fun{typeof(means)}(means) -CentralizedAbs2Fun(means, extrude) = CentralizedAbs2Fun{eltype(means)}(Broadcast.extrude(means)) +# An extruded means array pairs with `Base.PairsArray`: each element's mean is +# looked up by the element's index with broadcasting semantics +CentralizedAbs2Fun(means::Broadcast.Extruded) = CentralizedAbs2Fun{eltype(means.x)}(means) # Division is generally costly, but Julia is typically able to constant propagate a /1 # and simply ensure we get the type right at no cost, allowing the division in-place later (f::CentralizedAbs2Fun)(x) = abs2.(x - f.mean)/1 -(f::CentralizedAbs2Fun{<:Any,<:Broadcast.Extruded})((i, x),) = abs2.(x - Broadcast._broadcast_getindex(f.mean, i))/1 +function (f::CentralizedAbs2Fun{<:Any,<:Broadcast.Extruded})((i, x),) + # The means' axes are checked to match the kept axes of `A` before any + # reduction begins (in `varm!` and `_varm`), so the lookup is inbounds + j = Broadcast.newindex(i, f.mean.keeps, f.mean.defaults) + m = @inbounds f.mean.x[j] + return abs2.(x - m)/1 +end _doubled(x) = x+x Base.mapreduce_empty(::CentralizedAbs2Fun{T,<:Broadcast.Extruded}, ::typeof(Base.add_sum), ::Type{Tuple{_Any,S}}) where {T<:Number, S<:Number, _Any} = _doubled(abs2(zero(T)-zero(S)))/1 Base.mapreduce_empty(::CentralizedAbs2Fun{T,<:Broadcast.Extruded}, ::typeof(Base.add_sum), ::Type{Tuple{_Any, Union{Missing, S}}}) where {T<:Number, S<:Number, _Any} = _doubled(abs2(zero(T)-zero(S)))/1 @@ -277,9 +285,9 @@ Base.mapreduce_empty(::CentralizedAbs2Fun{T}, ::typeof(Base.add_sum), ::Type{Uni centralize_sumabs2(A::AbstractArray, m) = sum(CentralizedAbs2Fun(m), A) centralize_sumabs2(A::AbstractArray, m::AbstractArray, region) = - sum(CentralizedAbs2Fun(m, true), Base.PairsArray(A), dims=region) + sum(CentralizedAbs2Fun(Broadcast.extrude(m)), Base.PairsArray(A), dims=region) centralize_sumabs2!(R::AbstractArray, A::AbstractArray, means::AbstractArray) = - sum!(CentralizedAbs2Fun(means, true), R, Base.PairsArray(A)) + sum!(CentralizedAbs2Fun(Broadcast.extrude(means)), R, Base.PairsArray(A)) function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S @@ -346,7 +354,9 @@ function _varm(A::AbstractArray, m, corrected::Bool, region) if rn <= 0 R .= R ./ 0 else - R .= R .* 1//rn # why use Rational? + # Scaling by a Rational rounds only once for floating point eltypes + # and stays exact for Rational and Integer ones + R .= R .* (1 // rn) end return R end diff --git a/test/runtests.jl b/test/runtests.jl index 16761dba..cd0d0318 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1152,3 +1152,18 @@ end @test isequal(mean(x; dims=1), [2. 4. missing 8.]) @test isequal(mean(x; dims=2), [2.5; 5.0; missing;;]) end + +@testset "pairwise accuracy of var and mean" begin + # The reduction-based implementation is pairwise: accumulating the + # centralized squares of 10^6 Float32 values centered on 1f4 in a naive + # left-to-right loop loses three to four digits of the variance, while + # the pairwise reduction stays within a few eps + x = randn(MersenneTwister(1), Float32, 10^6) .+ 1f4 + v = var(Float64.(x)) + @test var(x) ≈ v rtol=1e-5 + @test var(reshape(x, :, 1); dims=1)[1] ≈ v rtol=1e-5 + @test var(reshape(x, 1, :); dims=2)[1] ≈ v rtol=1e-5 + m = mean(Float64.(x)) + @test mean(x) ≈ m rtol=1e-6 + @test varm(x, Float32(m)) ≈ v rtol=1e-5 +end From ba9787e53de30a5dc19ca5deaa949b1d0ae75186 Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Thu, 30 Jul 2026 23:30:53 -0400 Subject: [PATCH 5/6] test: align the suite with the pairwise reduction machinery Detect the new reduction machinery positively, via `isdefined(Base, :mapreduce_similar)`, instead of checking for a removed internal. Collapse the duplicated cov/cor missing-propagation branches into single `@test ... broken=` lines. Restore the exact `=== NaN` assertion for the variance of an empty matrix, since the whole-array path returns the canonical NaN again. Add coverage that was missing while porting: `mean`/`var`/`std` over general iterators (generators and `skipmissing`), and a check that `var` over every dims combination of a 3-d array matches a naive centralized-sum computation. Assisted-By: Claude Fable 5 --- test/runtests.jl | 58 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 39 insertions(+), 19 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index cd0d0318..1c6157b9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,6 +3,11 @@ using Statistics, Test, Random, LinearAlgebra, SparseArrays, Dates using Test: guardseed +# The pairwise reduction machinery (JuliaLang/julia#58418) supports +# heterogeneous eltypes in dimensional reductions, which lets cov/cor +# propagate `missing` instead of throwing +const NEW_REDUCTION_MACHINERY = isdefined(Base, :mapreduce_similar) + Random.seed!(123) @testset "middle" begin @@ -375,7 +380,7 @@ end @testset "var: empty cases" begin A = Matrix{Int}(undef, 0,1) - @test isnan(var(A)) + @test var(A) === NaN @test isequal(var(A, dims=1), fill(NaN, 1, 1)) @test isequal(var(A, dims=2), fill(NaN, 0, 1)) @@ -561,15 +566,9 @@ Y = [6.0 2.0; @testset "cov with missing" begin @test cov([missing]) === cov([1, missing]) === missing @test cov([1, missing], [2, 3]) === cov([1, 3], [2, missing]) === missing - if isdefined(Base, :_reducedim_init) - @test_broken isequal(coalesce.(cov([1 missing; 2 3]), NaN), cov([1 NaN; 2 3])) - @test_broken isequal(coalesce.(cov([1 missing; 2 3], [1, 2]), NaN), cov([1 NaN; 2 3], [1, 2])) - @test_broken isequal(coalesce.(cov([1, 2], [1 missing; 2 3]), NaN), cov([1, 2], [1 NaN; 2 3])) - else - @test isequal(coalesce.(cov([1 missing; 2 3]), NaN), cov([1 NaN; 2 3])) - @test isequal(coalesce.(cov([1 missing; 2 3], [1, 2]), NaN), cov([1 NaN; 2 3], [1, 2])) - @test isequal(coalesce.(cov([1, 2], [1 missing; 2 3]), NaN), cov([1, 2], [1 NaN; 2 3])) - end + @test isequal(coalesce.(cov([1 missing; 2 3]), NaN), cov([1 NaN; 2 3])) broken=!NEW_REDUCTION_MACHINERY + @test isequal(coalesce.(cov([1 missing; 2 3], [1, 2]), NaN), cov([1 NaN; 2 3], [1, 2])) broken=!NEW_REDUCTION_MACHINERY + @test isequal(coalesce.(cov([1, 2], [1 missing; 2 3]), NaN), cov([1, 2], [1 NaN; 2 3])) broken=!NEW_REDUCTION_MACHINERY @test isequal(cov([1 2; 2 3], [1, missing]), [missing missing]') @test isequal(cov([1, missing], [1 2; 2 3]), [missing missing]) end @@ -678,15 +677,9 @@ end @test cor([missing]) === missing @test cor([1, missing]) == 1 @test cor([1, missing], [2, 3]) === cor([1, 3], [2, missing]) === missing - if isdefined(Base, :_reducedim_init) - @test_broken isequal(coalesce.(cor([1 missing; 2 3]), NaN), cor([1 NaN; 2 3])) - @test_broken isequal(coalesce.(cor([1 missing; 2 3], [1, 2]), NaN), cor([1 NaN; 2 3], [1, 2])) - @test_broken isequal(coalesce.(cor([1, 2], [1 missing; 2 3]), NaN), cor([1, 2], [1 NaN; 2 3])) - else - @test isequal(coalesce.(cor([1 missing; 2 3]), NaN), cor([1 NaN; 2 3])) - @test isequal(coalesce.(cor([1 missing; 2 3], [1, 2]), NaN), cor([1 NaN; 2 3], [1, 2])) - @test isequal(coalesce.(cor([1, 2], [1 missing; 2 3]), NaN), cor([1, 2], [1 NaN; 2 3])) - end + @test isequal(coalesce.(cor([1 missing; 2 3]), NaN), cor([1 NaN; 2 3])) broken=!NEW_REDUCTION_MACHINERY + @test isequal(coalesce.(cor([1 missing; 2 3], [1, 2]), NaN), cor([1 NaN; 2 3], [1, 2])) broken=!NEW_REDUCTION_MACHINERY + @test isequal(coalesce.(cor([1, 2], [1 missing; 2 3]), NaN), cor([1, 2], [1 NaN; 2 3])) broken=!NEW_REDUCTION_MACHINERY @test isequal(cor([1 2; 2 3], [1, missing]), [missing missing]') @test isequal(cor([1, missing], [1 2; 2 3]), [missing missing]) end @@ -1167,3 +1160,30 @@ end @test mean(x) ≈ m rtol=1e-6 @test varm(x, Float32(m)) ≈ v rtol=1e-5 end + +@testset "mean and var of general iterators" begin + g = (x^2 for x in 1:4) + @test mean(g) === mean([1, 4, 9, 16]) === 7.5 + @test mean(sqrt, x^2 for x in 1:3) === mean([1.0, 2.0, 3.0]) + @test var(g) === var([1, 4, 9, 16]) + @test var(g; corrected=false) === var([1, 4, 9, 16]; corrected=false) + @test var(g; mean=7.5) === var([1, 4, 9, 16]; mean=7.5) + @test std(g) === std([1, 4, 9, 16]) + s = skipmissing([1.0, missing, 2.0, 3.0]) + @test mean(s) === 2.0 + @test var(s) === var([1.0, 2.0, 3.0]) + @test var(s; mean=2.0) === var([1.0, 2.0, 3.0]; mean=2.0) + @test varm(s, missing) === missing +end + +@testset "var over multiple dims matches the naive computation" begin + B = randn(MersenneTwister(3), 3, 4, 5) + for d in (1, 2, 3, (1, 2), (1, 3), (2, 3), (1, 2, 3)) + m = mean(B, dims=d) + n = prod(size(B, i) for i in d) + @test var(B; dims=d) ≈ sum(abs2, B .- m; dims=d) ./ (n - 1) + @test var(B; dims=d, corrected=false) ≈ sum(abs2, B .- m; dims=d) ./ n + @test var(B; dims=d, mean=m) ≈ var(B; dims=d) + @test std(B; dims=d) ≈ sqrt.(var(B; dims=d)) + end +end From a7ef77466b4dd9cffa6b2f4a31db5a8f6d2e4e01 Mon Sep 17 00:00:00 2001 From: Matt Bauman Date: Thu, 6 Aug 2026 18:08:43 -0400 Subject: [PATCH 6/6] more carefully ensure means bounds are checked This is the only place `CentralizedAbs2Fun{<:Extruded}` is constructed. Incidentally, the SparseArray extension already had its `centralize_sumabs2` function do this checking as well. It generally seems to fit better at this level (otherwise I would want to call it `unsafe_centralize_sumabs2`). --- src/Statistics.jl | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/Statistics.jl b/src/Statistics.jl index 3686af05..35a1b0e2 100644 --- a/src/Statistics.jl +++ b/src/Statistics.jl @@ -270,8 +270,7 @@ CentralizedAbs2Fun(means::Broadcast.Extruded) = CentralizedAbs2Fun{eltype(means. # and simply ensure we get the type right at no cost, allowing the division in-place later (f::CentralizedAbs2Fun)(x) = abs2.(x - f.mean)/1 function (f::CentralizedAbs2Fun{<:Any,<:Broadcast.Extruded})((i, x),) - # The means' axes are checked to match the kept axes of `A` before any - # reduction begins (in `varm!` and `_varm`), so the lookup is inbounds + # The means' axes are checked in `centralize_sumabs2` and `centralize_sumabs2!` j = Broadcast.newindex(i, f.mean.keeps, f.mean.defaults) m = @inbounds f.mean.x[j] return abs2.(x - m)/1 @@ -284,17 +283,16 @@ Base.mapreduce_empty(::CentralizedAbs2Fun{T}, ::typeof(Base.add_sum), ::Type{Uni centralize_sumabs2(A::AbstractArray, m) = sum(CentralizedAbs2Fun(m), A) -centralize_sumabs2(A::AbstractArray, m::AbstractArray, region) = - sum(CentralizedAbs2Fun(Broadcast.extrude(m)), Base.PairsArray(A), dims=region) -centralize_sumabs2!(R::AbstractArray, A::AbstractArray, means::AbstractArray) = - sum!(CentralizedAbs2Fun(Broadcast.extrude(means)), R, Base.PairsArray(A)) - +function centralize_sumabs2(A::AbstractArray, means::AbstractArray, region) + _checkm(A, means, region) + return sum(CentralizedAbs2Fun(Broadcast.extrude(means)), Base.PairsArray(A), dims=region) +end +function centralize_sumabs2!(R::AbstractArray, A::AbstractArray, means::AbstractArray) + _checkm(R, means) + return sum!(CentralizedAbs2Fun(Broadcast.extrude(means)), R, Base.PairsArray(A)) +end function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; corrected::Bool=true) where S - for d in 1:max(ndims(R), ndims(m)) - axes(m, d) == axes(R, d) || throw(DimensionMismatch( - "dimension $d of `mean` should have indices $(axes(R, d)), but got $(axes(m, d))")) - end if isempty(A) fill!(R, convert(S, NaN)) else @@ -337,7 +335,7 @@ over dimensions. In that case, `mean` must be an array with the same shape as """ varm(A::AbstractArray, m::AbstractArray; corrected::Bool=true, dims=:) = _varm(A, m, corrected, dims) -_throw_mean_mismatch(A, m, region) = throw(DimensionMismatch("axes of means ($(axes(m))) does not match reduction over $(region) of $(axes(A))")) +_throw_mean_mismatch(A, m, region) = throw(DimensionMismatch("axes of means ($(axes(m))) do not match reduction over $(region) of $(axes(A))")) function _checkm(A::AbstractArray, m::AbstractArray, region) for d in 1:max(ndims(A), ndims(m)) if d in region @@ -347,8 +345,14 @@ function _checkm(A::AbstractArray, m::AbstractArray, region) end end end +_throw_mean_mismatch(R, m) = throw(DimensionMismatch("axes of means ($(axes(m))) do not match reduction into $(axes(R))")) +function _checkm(R::AbstractArray, m::AbstractArray) + for d in 1:max(ndims(R), ndims(m)) + axes(m, d) == axes(R, d) || _throw_mean_mismatch(R, m) + end +end + function _varm(A::AbstractArray, m, corrected::Bool, region) - _checkm(A, m, region) rn = prod(ntuple(d->d in region ? size(A, d) : 1, Val(ndims(A)))) - Int(corrected) R = centralize_sumabs2(A, m, region) if rn <= 0