Proposal: generalise logsumexp slightly.

[willtebbutt]

Will add thorough testing if this is deemed something that's reasonable to include. Provides functionality similar to that found in sum and maximum etc to provide a dims argument which acts in the same way as sum and maximum.

e.g.

B = randn(5, 4)
A = logsumexp(B; dims=1)

returns a 1 x 4 matrix A, where A[1, k] = logsumexp(B[:, k]).

[ararslan]

I think it's more contentional to make this something like

Suggested change

	dims isa Nothing && return _logsumexp(X)
	dims === nothing && return _logsumexp(X)

though it doesn't really matter; they're entirely equivalent and should perform the same. Just figured I'd note it.

[cossio]

isnothing(dims)?

[ararslan]

=== is special cased in the compiler and is generally more efficient for checking nothing (and indeed, even missing).

[cossio]

I think isnothing is also elided since it's dealt with at dispatch time, and it was the recommended way to check for nothingness? The only thing is that it requires Julia 1.1.

[ararslan]

it was the recommended way to check for nothingness?

How did you come to that conclusion?

I think isnothing is also elided

It is not, see JuliaLang/julia#27681.

[cossio]

How did you come to that conclusion?

Well, it's a exported function named isnothing
In reality I wasn't aware of that issue. Thanks for pointing it out.

[ararslan]

This certainly looks like a worthwhile feature to me!

[andreasnoack]

Just two minor comments.

[andreasnoack]

Isn't this type unstable? I think you'd still have to consider dimensions as in

julia> sum(zeros(0, 0), dims=1)
1×0 Array{Float64,2}

[andreasnoack]

Maybe reuse the array created by sum(exp.(X .- u); dims=dims) to avoid a temporary.

[tpapp]

Is this something that cannot be addressed with mapslices and similar?

[willtebbutt]

Are there not performance issues associated with mapslices and related?

[andreasnoack]

Are there not performance issues associated with mapslices and related?

Unfortunately, there are. At least there used to be. Could you benchmark it?

[willtebbutt]

Here are some timings:

using BenchmarkTools, StatsFuns

N1, M1 = 1000, 2;
x = randn(N1, M1);
@btime logsumexp($x);
  20.012 μs (0 allocations: 0 bytes)
@btime logsumexp($x; dims=1);
  25.779 μs (13 allocations: 16.23 KiB)
@btime mapslices(logsumexp, $x; dims=1);
  25.318 μs (56 allocations: 10.53 KiB)
@btime logsumexp($x; dims=2);
  33.160 μs (31 allocations: 40.33 KiB)
@btime mapslices(logsumexp, $x; dims=2);
  401.447 μs (9507 allocations: 236.17 KiB)

# Second round of tests with different memory layout.
x2 = Matrix(x');
@btime logsumexp($x2);
  19.984 μs (0 allocations: 0 bytes)
@btime logsumexp($x2; dims=2);
  31.323 μs (31 allocations: 16.80 KiB)
@btime mapslices(logsumexp, $x2; dims=2);
  25.503 μs (56 allocations: 10.53 KiB)
@btime logsumexp($x2; dims=1);
  38.780 μs (13 allocations: 39.77 KiB)
@btime mapslices(logsumexp, $x2; dims=1);
  373.496 μs (9507 allocations: 236.17 KiB)

When each slice is small the mapslices overhead associated with mapslices is rather large.

[tpapp]

Thanks for the benchmarks. This is because mapslices is poorly implemented. Cf

using BenchmarkTools, StatsFuns, JuliennedArrays

N1, M1 = 1000, 2;
x = randn(N1, M1);
@btime mapslices(logsumexp, $x; dims=2);
@btime map(logsumexp, Slices($x, False(), True()));

with

julia> @btime mapslices(logsumexp, $x; dims=2);
  704.918 μs (9507 allocations: 236.17 KiB)

julia> @btime map(logsumexp, Slices($x, False(), True()));
  52.297 μs (1003 allocations: 54.84 KiB)

I think it would be better to propagate the use of sane slice iterations constructs instead of adding a (; dims = ...) method to everything.

[willtebbutt]

Closing for now as this has gone stale and a more generic mechanism of the form @tpapp suggests would probably better. Happy to re-open if anyone feels strongly about this.