๏‚›Edit on GitHub

Update: Should probably use StatsFuns.jl instead

This function exists in StatsFuns.jl (wish I had found that on my first search...). Their implementation is fancier and a little more efficient. Recommend going with theirs unless you really, really want to keep your dependencies small.

LogSumExp.jl

A Julia implementation of logsumexp.

LogSumExp.logsumexp โ€” Method
logsumexp(a; dims=:)

Computes the log-sum-exp. Mathematically equivalent to

log.(sum(exp.(a); dims=dims))

but more numerically stable.

Note on treatment of dims

Note that, in keeping with the other Julia array operations, no dims are dropped in the output of logsumexp; instead,

size(logsumexp(a; dims=dims), d) == 1

for each d in dims. For example:

julia> a = rand(2, 3, 4, 5);

julia> size(logsumexp(a; dims=(1,3)))
(1, 3, 1, 5)

For a variant of logsumexp that does drop dims, see reduce_logsumexp.

source
LogSumExp.reduce_logsumexp โ€” Method
reduce_logsumexp(a; dims=:)

Like logsumexp, but drops the dims dims in the output.

julia> a = rand(2, 3, 4, 5);

julia> size(reduce_logsumexp(a; dims=(1,3)))
(3, 5)
source