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,90 @@ 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.
+`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)
+
+"""
+ 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)
 """
-function logsumexp(X)
+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
+
+# initial element is required by CUDA (ideally we would never use this method)
+function _logsumexp_onepass(X, ::Base.HasEltype)
+ # do not perform type computations if element type is abstract
+ T = eltype(X)
+ isconcretetype(T) || return _logsumexp_onepass(X, Base.EltypeUnknown())
+
+ # compute initial element
+ FT = float(eltype(X))
+ init = (FT(-Inf), zero(FT))
+ r_one = one(FT)
+
+ return mapreduce(_logsumexp_onepass_op, X; init=init) do x
+ return float(x), r_one
  end
 end
 
+# without initial element (ideally we would always use this method)
+function _logsumexp_onepass(X, ::Base.EltypeUnknown)
+ return mapreduce(_logsumexp_onepass_op, X) do x
+ _x = float(x)
+ return _x, one(_x)
+ end
+end
+
+# all inputs are provided as floating point numbers
+# `r1` and `r2` are one if `xmax1` and `xmax2` are new elements (no partial sums)
+function _logsumexp_onepass_op((xmax1, r1)::T, (xmax2, r2)::T) where {T<:Tuple}
+ if xmax1 < xmax2
+ xmax = xmax2
+ a = exp(xmax1 - xmax2)
+ r = r2 + (isone(r1) ? a : r1 * a) # avoid expensive multiplication for new elements
+ elseif xmax1 > xmax2
+ xmax = xmax1
+ a = exp(xmax2 - xmax1)
+ r = r1 + (isone(r2) ? a : r2 * a) # avoid expensive multiplication for new elements
+ else # ensure finite values if x = xmax = ± Inf
+ xmax = isnan(xmax1) ? xmax1 : xmax2
+ r = r1 + r2
+ end
+
+ return xmax, r
+end
 
 """
  softmax!(r::AbstractArray, x::AbstractArray)

test/basicfuns.jl

@@ -137,6 +137,10 @@ end
  @test isnan(logsumexp([NaN, 9.0]))
  @test isnan(logsumexp([NaN, Inf]))
  @test isnan(logsumexp([NaN, -Inf]))
+
+ # logsumexp with general iterables (issue #63)
+ xs = range(-500, stop = 10, length = 1000)
+ @test logsumexp(x for x in xs) == logsumexp(xs)
 end
 
 @testset "softmax" begin