Fix logsumexp! with output arrays of abstract eltype (#40)

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.9"
+version = "0.3.10"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/logsumexp.jl

@@ -42,8 +42,8 @@ See also [`logsumexp`](@ref).
 
 [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)
+function logsumexp!(out::AbstractArray, X::AbstractArray{<:Number})
+    FT = float(eltype(X))
     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))

test/basicfuns.jl

@@ -298,6 +298,13 @@ end
     @test @inferred(logsumexp(xs; dims=2)) ≈ log.(sum(exp.(xs); dims=2))
     @test @inferred(logsumexp(xs; dims=[1, 2])) ≈ log(sum(exp.(xs); dims=[1, 2]))
     @test @inferred(logsumexp(x for x in xs)) == logsumexp(xs)
+
+    # output arrays with abstract eltype
+    xs = randn(2, 4)
+    out = [missing, 1.0]
+    expected = logsumexp(xs; dims=2)
+    @test logsumexp!(out, xs) ≈ expected
+    @test out ≈ expected
 end
 
 @testset "softmax" begin