Add logsumexp! (#39)

* Add `logsumexp!`

* Add some tests

* Update docstring

* Fix test error

* Fix test

* More test fixes

* Add more tests

* Update docstring

* Update docstrings

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.3.8"
+version = "0.3.9"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/index.md

@@ -25,6 +25,7 @@ logmxp1
 logaddexp
 logsubexp
 logsumexp
+logsumexp!
 softmax!
 softmax
 ```

src/LogExpFunctions.jl

@@ -10,7 +10,7 @@ import IrrationalConstants
 import LinearAlgebra
 
 export xlogx, xlogy, xlog1py, xexpx, xexpy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
-    softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, softmax,
+    softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, logsumexp!, softmax,
     softmax!, logcosh
 
 include("basicfuns.jl")

src/logsumexp.jl

@@ -1,11 +1,12 @@
 """
 $(SIGNATURES)
 
-Compute `log(sum(exp, X))` in a numerically stable way that avoids intermediate over- and
-underflow.
+Compute `log(sum(exp, X))`.
 
-`X` should be an iterator of real or complex numbers. The result is computed using a single
-pass over the data.
+`X` should be an iterator of real or complex numbers.
+The result is computed in a numerically stable way that avoids intermediate over- and underflow, using a single pass over the data.
+
+See also [`logsumexp!`](@ref).
 
 # References
 
@@ -16,17 +17,38 @@ logsumexp(X) = _logsumexp_onepass(X)
 """
 $(SIGNATURES)
 
-Compute `log.(sum(exp.(X); dims=dims))` in a numerically stable way that avoids
-intermediate over- and underflow.
+Compute `log.(sum(exp.(X); dims=dims))`.
+
+The result is computed in a numerically stable way that avoids intermediate over- and underflow, using a single pass over the data.
 
-The result is computed using a single pass over the data.
+See also [`logsumexp!`](@ref).
 
 # References
 
 [Sebastian Nowozin: Streaming Log-sum-exp Computation](http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html)
 """
 logsumexp(X::AbstractArray{<:Number}; dims=:) = _logsumexp(X, dims)
 
+"""
+$(SIGNATURES)
+
+Compute [`logsumexp`](@ref) of `X` over the singleton dimensions of `out`, and write results to `out`.
+
+The result is computed in a numerically stable way that avoids intermediate over- and underflow, using a single pass over the data.
+
+See also [`logsumexp`](@ref).
+
+# References
+
+[Sebastian Nowozin: Streaming Log-sum-exp Computation](http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html)
+"""
+function logsumexp!(out::AbstractArray{<:Number}, X::AbstractArray{<:Number})
+    FT = eltype(out)
+    xmax_r = fill!(similar(out, Tuple{FT,FT}), (FT(-Inf), zero(FT)))
+    Base.reducedim!(_logsumexp_onepass_op, xmax_r, X)
+    return @. out = first(xmax_r) + log1p(last(xmax_r))
+end
+
 _logsumexp(X::AbstractArray{<:Number}, ::Colon) = _logsumexp_onepass(X)
 function _logsumexp(X::AbstractArray{<:Number}, dims)
     # Do not use log(zero(eltype(X))) directly to avoid issues with ForwardDiff (#82)

test/basicfuns.jl

@@ -195,15 +195,31 @@ end
 
     _x = [[1.0, 2.0, 3.0] [1.0, 2.0, 3.0] .+ 1000.]
     for x in (_x, complex(_x))
-        @test @inferred(logsumexp(x; dims=1)) ≈ [3.40760596444438 1003.40760596444438]
-        @test @inferred(logsumexp(x; dims=[1, 2])) ≈ [1003.4076059644444]
+        expected = [3.40760596444438 1003.40760596444438]
+        @test @inferred(logsumexp(x; dims=1)) ≈ expected
+        out = Array{eltype(x)}(undef, 1, 2)
+        @test @inferred(logsumexp!(out, x)) ≈ expected
+        @test out ≈ expected
+
         y = copy(x')
-        @test @inferred(logsumexp(y; dims=2)) ≈ [3.40760596444438, 1003.40760596444438]
+        expected = [3.40760596444438, 1003.40760596444438]
+        @test @inferred(logsumexp(y; dims=2)) ≈ expected
+        out = Array{eltype(y)}(undef, 2)
+        @test @inferred(logsumexp!(out, y)) ≈ expected
+        @test out ≈ expected
+
+        expected = [1003.4076059644444]
+        @test @inferred(logsumexp(x; dims=[1, 2])) ≈ expected
+        out = Array{eltype(x)}(undef, 1)
+        @test @inferred(logsumexp!(out, x)) ≈ expected
+        @test out ≈ expected
     end
 
     # check underflow
     @test logsumexp([1e-20, log(1e-20)]) ≈ 2e-20
     @test logsumexp(Complex{Float64}[1e-20, log(1e-20)]) ≈ 2e-20
+    @test logsumexp!([1.0], [1e-20, log(1e-20)]) ≈ [2e-20]
+    @test logsumexp!(Complex{Float64}[1.0], Complex{Float64}[1e-20, log(1e-20)]) ≈ [2e-20]
 
     let cases = [([-Inf, -Inf], -Inf),   # correct handling of all -Inf
                  ([-Inf, -Inf32], -Inf), # promotion
@@ -216,6 +232,15 @@ end
             @test logaddexp(arguments...) ≡ result
             @test logsumexp(arguments) ≡ result
             @test logsumexp(complex(arguments)) ≡ complex(result)
+
+            FT = float(eltype(arguments))
+            out = [one(FT)]
+            @test logsumexp!(out, arguments)[1] ≡ result
+            @test out[1] ≡ result
+
+            out = [one(complex(FT))]
+            @test logsumexp!(out, complex(arguments))[1] ≡ complex(result)
+            @test out[1] ≡ complex(result)
         end
     end
 
@@ -250,6 +275,15 @@ end
     @test isnan(logsumexp(Complex{Float64}[NaN * im, 9.0]))
     @test isnan(logsumexp(Complex{Float64}[NaN * im, Inf]))
     @test isnan(logsumexp(Complex{Float64}[NaN * im, -Inf]))
+    @test isnan(logsumexp!([1.0], [NaN, 9.0])[1])
+    @test isnan(logsumexp!([1.0], [NaN, Inf])[1])
+    @test isnan(logsumexp!([1.0], [NaN, -Inf])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN, 9.0])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN, Inf])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN, -Inf])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, 9.0])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, Inf])[1])
+    @test isnan(logsumexp!(Complex{Float64}[1.0], Complex{Float64}[NaN * im, -Inf])[1])
 
     # logsumexp with general iterables (issue #63)
     xs = range(-500, stop = 10, length = 1000)