Skip to content

WIP: implement reductions as reductions #172

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mbauman wants to merge 7 commits into
base: master
Choose a base branch
from
Open

WIP: implement reductions as reductions #172

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6221324
WIP: implement reductions as reductions
mbauman Aug 27, 2024
c674ab8
skip the mean refactor for now
mbauman May 14, 2025
dba8482
Merge upstream master at e14c8f8 into mb/reduce-reuse-recycle
mbauman Jul 31, 2026
257e827
varm!: fix the shape check for array means and single-element behavior
mbauman Jul 31, 2026
9a80ac7
variance: index the extruded means directly with inbounds
mbauman Jul 31, 2026
ba9787e
test: align the suite with the pairwise reduction machinery
mbauman Jul 31, 2026
a7ef774
more carefully ensure means bounds are checked
mbauman Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 69 additions & 51 deletions src/Statistics.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ realXcY(x::Number, y::Number) = real(conj(x) * 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)
Expand Down Expand Up @@ -257,63 +258,53 @@ function _var(iterable, corrected::Bool, mean)
end
end

centralizedabs2fun(m) = x -> abs2.(x - m)
struct CentralizedAbs2Fun{T,S} <: Function
mean::S
end
CentralizedAbs2Fun{T}(means) where {T} = CentralizedAbs2Fun{T,typeof(means)}(means)
CentralizedAbs2Fun(means) = CentralizedAbs2Fun{typeof(means)}(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
function (f::CentralizedAbs2Fun{<:Any,<:Broadcast.Extruded})((i, x),)
# 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
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
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) =
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
sum(CentralizedAbs2Fun(m), 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
if isempty(A)
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)
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
Expand Down Expand Up @@ -344,8 +335,35 @@ 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))) 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
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
_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)
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
# 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

varm(A::AbstractArray, m; corrected::Bool=true) = _varm(A, m, corrected, :)

Expand Down
95 changes: 89 additions & 6 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -383,6 +388,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]

Expand Down Expand Up @@ -541,9 +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
@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])
@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
Expand Down Expand Up @@ -652,9 +677,9 @@ 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])
@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
Expand Down Expand Up @@ -1104,3 +1129,61 @@ 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_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))
@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

@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