[Python] Why is manual (-ish) computation of log-sum-exp faster than NumPy/SciPy functions?

I’m working with a very large array and was surprised to find that manually computing the log of the sum of exponentials is actually faster than using the built-in functions for this task from SciPy (scipy.special.logsumexp) and even NumPy (numpy.logaddexp.reduce)

I assumed the library implementations would be more optimized.. Am I doing something silly or missing something obvious?

# test_logsumexp.py

from functools import partial
import timeit
import numpy as np
from scipy.special import logsumexp


def logsumexp_manual(a, axis=-1):
max_vals = np.max(a, axis=axis)
exp_terms = np.exp(a - max_vals[..., np.newaxis])
sum_exp_terms = np.sum(exp_terms, axis=axis)
return max_vals + np.log(sum_exp_terms)

# Validity check
arr = np.random.rand(10_000, 30, 3)
scipy_result = logsumexp(arr, axis=-1)
numpy_result = np.logaddexp.reduce(arr, axis=-1)
manual_result = logsumexp_manual(arr)
np.testing.assert_allclose(scipy_result, manual_result)
np.testing.assert_allclose(scipy_result, numpy_result)


n_loops = 10
setup = ("import numpy as np; "
"from scipy.special import logsumexp; "
"arr = np.random.rand(10_000, 30, 3)")

total_scipy = timeit.timeit(
"logsumexp(arr, axis=-1)",
setup=setup,
number=n_loops,
)

total_numpy = timeit.timeit(
"np.logaddexp.reduce(arr, axis=-1)",
setup=setup,
number=n_loops,
)

logsumexp_partial = partial(logsumexp, a=arr)
total_manual = timeit.timeit(
logsumexp_partial,
setup=setup,
number=n_loops,
)

print(f"Scipy logsumexp: {total_scipy / n_loops:.6f} seconds per loop")
print(f"Numpy logaddexp: {total_numpy / n_loops:.6f} seconds per loop")
print(f"Manual logsumexp: {total_manual / n_loops:.6f} seconds per loop")

console output


Scipy logsumexp: 0.054794 seconds per loop
Numpy logaddexp: 0.032316 seconds per loop
Manual logsumexp: 0.010943 seconds per loop

Continue reading...