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,96 @@ 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.
+
+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)
+ FT = float(eltype(X))
+ xmax_r = reduce(_logsumexp_onepass_op, X; dims=dims, init=(FT(-Inf), zero(FT)))
+ return @. first(xmax_r) + log1p(last(xmax_r))
+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 + log1p(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
- end
+
+# initial element is required by CUDA (otherwise we could remove 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())
+
+ FT = float(T)
+ return reduce(_logsumexp_onepass_op, X; init=(FT(-Inf), zero(FT)))
+end
+
+# without initial element (without CUDA support we could always use this method)
+_logsumexp_onepass(X, ::Base.EltypeUnknown)::Tuple = reduce(_logsumexp_onepass_op, X)
+
+## Reductions for one-pass algorithm: avoid expensive multiplications if numbers are reduced
+
+# reduce two numbers
+function _logsumexp_onepass_op(x1, x2)
+ a = x1 == x2 ? zero(x1 - x2) : -abs(x1 - x2)
+ xmax = x1 > x2 ? oftype(a, x1) : oftype(a, x2)
+ r = exp(a)
+ return xmax, r
 end
 
+# reduce a number and a partial sum
+function _logsumexp_onepass_op(x, (xmax, r)::Tuple)
+ a = x == xmax ? zero(x - xmax) : -abs(x - xmax)
+ if x > xmax
+ _xmax = oftype(a, x)
+ _r = (r + one(r)) * exp(a)
+ else
+ _xmax = oftype(a, xmax)
+ _r = r + exp(a)
+ end
+ return _xmax, _r
+end
+_logsumexp_onepass_op(xmax_r::Tuple, x) = _logsumexp_onepass_op(x, xmax_r)
+
+# reduce two partial sums
+function _logsumexp_onepass_op((xmax1, r1)::Tuple, (xmax2, r2)::Tuple)
+ a = xmax1 == xmax2 ? zero(xmax1 - xmax2) : -abs(xmax1 - xmax2)
+ if xmax1 > xmax2
+ xmax = oftype(a, xmax1)
+ r = r1 + (r2 + one(r2)) * exp(a)
+ else
+ xmax = oftype(a, xmax2)
+ r = r2 + (r1 + one(r1)) * exp(a)
+ 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