Merge pull request #244 from FluxML/logsumexp

add logsumexp

Project.toml

@@ -16,8 +16,9 @@ Requires = "0.5, 1.0"
 julia = "1.3"
 
 [extras]
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Zygote"]
+test = ["Test", "Statistics", "Zygote"]

src/softmax.jl

@@ -1,5 +1,6 @@
 export softmax, softmax!, ∇softmax, ∇softmax!,
-       logsoftmax, logsoftmax!, ∇logsoftmax, ∇logsoftmax!
+       logsoftmax, logsoftmax!, ∇logsoftmax, ∇logsoftmax!,
+       logsumexp
 
 """
     softmax(x; dims=1)
@@ -114,3 +115,17 @@ end
 
 ∇logsoftmax(Δ, xs; dims=1) = Δ .- sum(Δ, dims=dims) .* softmax(xs, dims=dims)
 ∇logsoftmax!(Δ, xs) = ∇logsoftmax!(Δ, Δ, xs)
+
+"""
+    logsumexp(x; dims=:)
+
+Computes `log.(sum(exp.(x); dims=dims))` in a numerically stable
+way.
+
+See also [`logsoftmax`](@ref).
+"""
+function logsumexp(xs::AbstractArray; dims=:)
+    max_ = maximum(xs, dims=dims)
+    log_ = log.(sum(exp.(xs .- max_), dims=dims))
+    return max_ .+ log_
+end