Add onepass algorithm for logsumexp

src/basicfuns.jl

@@ -207,7 +207,7 @@ end
 Return `log(exp(x) + exp(y))`, avoiding intermediate overflow/undeflow, and handling non-finite values.
 """
 function logaddexp(x::Real, y::Real)
- # ensure Δ = 0 if x = y = Inf
+ # ensure Δ = 0 if x = y = ± Inf
  Δ = ifelse(x == y, zero(x - y), abs(x - y))
  max(x, y) + log1pexp(-Δ)
 end
@@ -224,28 +224,85 @@ logsubexp(x::Real, y::Real) = max(x, y) + log1mexp(-abs(x - y))
 """
  logsumexp(X)
 
-Compute `log(sum(exp, X))`, evaluated avoiding intermediate overflow/undeflow.
+Compute `log(sum(exp, X))` in a numerically stable way that avoids intermediate over- and
+underflow.
+
+`X` should be an iterator of real numbers. The result is computed using a single pass over
+the data.
+
+# References
+
+[Sebastian Nowozin: Streaming Log-sum-exp Computation.](http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html)
+"""
+logsumexp(X) = logsumexp_onepass(X)
 
-`X` should be an iterator of real numbers.
 """
-function logsumexp(X)
+ logsumexp(X::AbstractArray{<:Real}[; dims=:])
+
+Compute `log.(sum(exp.(X); dims=dims))` in a numerically stable way that avoids
+intermediate over- and underflow.
+
+If `dims = :`, then the result is computed using a single pass over the data.
+
+# References
+
+[Sebastian Nowozin: Streaming Log-sum-exp Computation.](http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html)
+"""
+logsumexp(X::AbstractArray{<:Real}; dims=:) = _logsumexp(X, dims)
+
+_logsumexp(X::AbstractArray{<:Real}, ::Colon) = logsumexp_onepass(X)
+function _logsumexp(X::AbstractArray{<:Real}, dims)
+ # Do not use log(zero(eltype(X))) directly to avoid issues with ForwardDiff (#82)
+ u = reduce(max, X, dims=dims, init=oftype(log(zero(eltype(X))), -Inf))
+ return u .+ log.(sum(exp.(X .- u); dims=dims))
+end
+
+function logsumexp_onepass(X)
+ # fallback for empty collections
  isempty(X) && return log(sum(X))
- reduce(logaddexp, X)
+
+ xmax, r = _logsumexp_onepass(X, Base.IteratorEltype(X))
+
+ return xmax + log(r)
 end
-function logsumexp(X::AbstractArray{T}; dims=:) where {T<:Real}
- # Do not use log(zero(T)) directly to avoid issues with ForwardDiff (#82)
- u = reduce(max, X, dims=dims, init=oftype(log(zero(T)), -Inf))
- u isa AbstractArray || isfinite(u) || return float(u)
- let u=u # avoid https://github.com/JuliaLang/julia/issues/15276
-  # TODO: remove the branch when JuliaLang/julia#31020 is merged.
-  if u isa AbstractArray
- u .+ log.(sum(exp.(X .- u); dims=dims))
-  else
-  u + log(sum(x -> exp(x-u), X))
- end
+
+# with initial element: required by CUDA
+function _logsumexp_onepass(X, ::Base.HasEltype)
+ # compute initial element
+ FT = float(eltype(X))
+ init = (FT(-Inf), zero(FT))
+ r_one = one(FT)
+
+ # perform single pass over the data
+ return mapreduce(_logsumexp_onepass_op, X; init=init) do x
+ return float(x), r_one
  end
 end
 
+# without initial element
+function _logsumexp_onepass(X, ::Base.EltypeUnknown)
+ return mapreduce(_logsumexp_onepass_op, X) do x
+ _x = float(x)
+ return _x, one(_x)
+ end
+end
+
+function _logsumexp_onepass_op((xmax1, r1)::T, (xmax2, r2)::T) where {T<:Tuple}
+ if xmax1 < xmax2
+ xmax = xmax2
+ a = exp(xmax1 - xmax2)
+ r = r2 + ifelse(isone(r1), a, r1 * a)
+ elseif xmax1 > xmax2
+ xmax = xmax1
+ a = exp(xmax2 - xmax1)
+ r = r1 + ifelse(isone(r2), a, r2 * a)
+ else # ensure finite values if x = xmax = ± Inf
+ xmax = ifelse(isnan(xmax1), xmax1, xmax2)
+ r = r1 + r2
+ end
+
+ return xmax, r
+end
 
 """
  softmax!(r::AbstractArray, x::AbstractArray)

test/basicfuns.jl

@@ -1,4 +1,5 @@
 using StatsFuns, Test
+using StatsFuns: logsumexp_onepass
 
 @testset "xlogx & xlogy" begin
  @test iszero(xlogx(0))
@@ -137,6 +138,11 @@ end
  @test isnan(logsumexp([NaN, 9.0]))
  @test isnan(logsumexp([NaN, Inf]))
  @test isnan(logsumexp([NaN, -Inf]))
+
+ # issue #63
+ a = logsumexp(i for i in range(-500, stop = 10, length = 1000) if true)
+ b = logsumexp(range(-500, stop = 10, length = 1000))
+ @test a == b
 end
 
 @testset "softmax" begin