Add onepass algorithm for logsumexp

src/basicfuns.jl

@@ -227,46 +227,40 @@ logsubexp(x::Real, y::Real) = max(x, y) + log1mexp(-abs(x - y))
 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.
 
-See also: [`logsumexp_onepass`](@ref)
+# 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=:) = _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))
-    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
-end
 
 """
-    logsumexp_onepass(X)
+    logsumexp(X::AbstractArray{<:Real}[; dims=:])
 
-Compute `log(sum(exp, X))` in a numerically stable way that avoids intermediate under- and
-overflow.
+Compute `log.(sum(exp.(X); dims=dims))` in a numerically stable way that avoids
+intermediate over- and underflow.
 
-In contrast to [`logsumexp`](@ref) the result is computed using a single pass over the data.
-`X` should be an iterator of real numbers.
+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))
 
-    # perform single pass over the data
     xmax, r = _logsumexp_onepass(X, Base.IteratorEltype(X))
 
     return xmax + log(r)

test/basicfuns.jl

@@ -98,10 +98,8 @@ end
     @test logaddexp(10002, 10003) ≈ 10000 + logaddexp(2.0, 3.0)
 
     @test logsumexp([1.0, 2.0, 3.0])          ≈ 3.40760596444438
-    @test logsumexp_onepass([1.0, 2.0, 3.0])  ≈ 3.40760596444438
     @test logsumexp((1.0, 2.0, 3.0))          ≈ 3.40760596444438
     @test logsumexp([1.0, 2.0, 3.0] .+ 1000.) ≈ 1003.40760596444438
-    @test logsumexp_onepass([1.0, 2.0, 3.0] .+ 1000.) ≈ 1003.40760596444438
 
     @test logsumexp([[1.0, 2.0, 3.0] [1.0, 2.0, 3.0] .+ 1000.]; dims=1) ≈ [3.40760596444438 1003.40760596444438]
     @test logsumexp([[1.0 2.0 3.0]; [1.0 2.0 3.0] .+ 1000.]; dims=2) ≈ [3.40760596444438, 1003.40760596444438]
@@ -117,7 +115,6 @@ end
         for (arguments, result) in cases
             @test logaddexp(arguments...) ≡ result
             @test logsumexp(arguments) ≡ result
-            @test logsumexp_onepass(arguments) ≡ result
         end
     end
 
@@ -142,10 +139,6 @@ end
     @test isnan(logsumexp([NaN, Inf]))
     @test isnan(logsumexp([NaN, -Inf]))
 
-    @test isnan(logsumexp_onepass([NaN, 9.0]))
-    @test isnan(logsumexp_onepass([NaN, Inf]))
-    @test isnan(logsumexp_onepass([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))