From 78bc9a276fd391b8f493359aff4ca2b7387d5c2a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 19 Oct 2026 18:43:40 +0200 Subject: [PATCH] Accumulate `mean` in floating point and divide afterwards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summing integers before dividing overflowed for the whole-array mean (`mean(fill(typemax(Int)÷2, 4))` gave -1.0), while the `dims` variant pre-scaled every term by 1/n, which loses precision (a mean of ones that is not one) and captured `f` in a `ComposedFunction`, so a type like `mean(Float32, A; dims=1)` could not be passed to the kernel. Follow Statistics.jl instead: convert each mapped element to the result's floating-point type, sum, then divide. The type is inferred rather than taken from `first(A)` to avoid scalar indexing, and the conversion closure keeps `f` as a static parameter when it is a type, so it stays isbits. The two-argument `_mean` methods are dropped: Statistics has only called `_mean(f, A, dims)` since 1.9. Fixes #819, fixes #611, supersedes #453. --- src/host/statistics.jl | 29 +++++++++++++++++------------ test/testsuite/statistics.jl | 23 +++++++++++++++++++++++ 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/src/host/statistics.jl b/src/host/statistics.jl index 3c6b0cc6c..5940a27b8 100644 --- a/src/host/statistics.jl +++ b/src/host/statistics.jl @@ -22,21 +22,26 @@ Statistics._std(A::AbstractGPUArray, corrected::Bool, mean, dims) = Statistics._std(A::AbstractGPUArray, corrected::Bool, mean, ::Colon) = sqrt.(Statistics.var(A; corrected=corrected, mean=mean, dims=:)) -# Revert https://github.com/JuliaLang/Statistics.jl/pull/25 -Statistics._mean(A::AbstractGPUArray, ::Colon) = sum(A) / length(A) -Statistics._mean(f, A::AbstractGPUArray, ::Colon) = sum(f, A) / length(A) - -function Statistics._mean(A::AbstractGPUArray, dims) - T = float(eltype(A)) - λ = convert(T, inv(_mean_denom(A, dims))) - sum(Base.Fix1(*,λ), A; dims) -end +# Like Statistics.jl, accumulate in the floating-point type of the result and divide +# afterwards: summing integers overflows, and scaling each term first loses precision. +# The accumulator type is inferred, as computing it from `first(A)` would need scalar indexing. function Statistics._mean(f, A::AbstractGPUArray, dims) - T = float(eltype(A)) - λ = convert(T, inv(_mean_denom(A, dims))) - sum(Base.Fix1(*,λ) ∘ f, A; dims) + isempty(A) && return sum(f, A; dims) / 0 + S = Base.promote_op(f, eltype(A)) + T = promote_type(Base.promote_op(/, S, Int), S) + result = sum(_mean_convert(f, T), A; dims) + if dims isa Colon + result / length(A) + else + result ./= _mean_denom(A, dims) + end end +# `f` and `T` are static parameters rather than captured values, so the closure stays isbits +# when `f` is a type (e.g. `mean(Float32, A)`) and can be passed to a kernel. +_mean_convert(f::F, ::Type{T}) where {F,T} = x -> convert(T, f(x)) +_mean_convert(::Type{F}, ::Type{T}) where {F,T} = x -> convert(T, F(x)) + function Statistics.covzm(x::AbstractGPUMatrix, vardim::Int=1; corrected::Bool=true) C = Statistics.unscaled_covzm(x, vardim) T = promote_type(typeof(one(eltype(C)) / 1), eltype(C)) diff --git a/test/testsuite/statistics.jl b/test/testsuite/statistics.jl index 29886163c..e453bcb21 100644 --- a/test/testsuite/statistics.jl +++ b/test/testsuite/statistics.jl @@ -38,6 +38,29 @@ using Statistics end end + @testset "mean accumulation" begin + # integers are summed in floating point, so they don't overflow + if Int64 in eltypes && Float64 in eltypes + x = fill(typemax(Int64) ÷ 2, 4, 2) + @test compare(mean, AT, x) + @test compare(A -> mean(A; dims = 1), AT, x) + @test compare(A -> mean(identity, A), AT, x) + @test compare(A -> mean(identity, A; dims = 2), AT, x) + @test compare(mean, AT, zeros(Int64, 0, 2)) nans = true + @test compare(A -> mean(A; dims = 1), AT, zeros(Int64, 0, 2)) nans = true + end + if Int32 in eltypes && Float32 in eltypes + x = rand(Int32, 10, 3) + @test compare(A -> mean(Float32, A), AT, x) + @test compare(A -> mean(Float32, A; dims = 1), AT, x) + end + + # dividing after summing is exact here, scaling every term first is not + if Float32 in eltypes + @test all(isone, Array(mean(AT(ones(Float32, 640, 64)); dims = 1))) + end + end + for ET in eltypes # Doesn't work with ComplexF32 in oneAPI for some reason. if !(ET in [Float32, Float64, Float16, ComplexF64])