Skip to content

Commit

Permalink
Call user function only once in mean
Browse files Browse the repository at this point in the history
Use an internal type to initiate promotion instead of using the type
returned by user-provided function applied to the first element.
This avoid calling the function twice, which can be confusing.
  • Loading branch information
nalimilan committed Sep 9, 2023
1 parent 81a90af commit ac9937d
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/Statistics.jl
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,10 @@ if !isdefined(Base, :mean)
"""
mean(A::AbstractArray; dims=:) = _mean(identity, A, dims)

struct _InitType end

_mean_promote(x::T, y::S) where {T,S} = convert(promote_type(T, S), y)
_mean_promote(x::_InitType, y::Any) = y/1

# ::Dims is there to force specializing on Colon (as it is a Function)
function _mean(f, A::AbstractArray, dims::Dims=:) where Dims
Expand All @@ -188,8 +191,7 @@ if !isdefined(Base, :mean)
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)
result = sum(x -> _mean_promote(_InitType(), f(x)), A, dims=dims)
if dims === (:)
return result / n
else
Expand Down

0 comments on commit ac9937d

Please sign in to comment.