Add `logsumexp`

aesara/tensor/math.py

@@ -2764,6 +2764,34 @@ def power(x, y):
     return x ** y
 
 
+def logsumexp(x, axis=None, keepdims=False):
+    """Compute the log of the sum of exponentials of input elements.
+
+    See ``scipy.special.logsumexp``.
+
+    Parameters
+    ----------
+    x : symbolic tensor
+        Input
+
+    axis : None or int or tuple of ints, optional
+        Axis or axes over which the sum is taken. By default axis is None,
+        and all elements are summed.
+
+    keepdims : bool, optional
+        If this is set to True, the axes which are reduced are left in the
+        result as dimensions with size one. With this option, the result will
+        broadcast correctly against the original array.
+
+    Returns
+    -------
+    tensor
+
+    """
+
+    return log(sum(exp(x), axis=axis, keepdims=keepdims))
+
+
 __all__ = [
     "max_and_argmax",
     "max",
@@ -2885,4 +2913,5 @@ __all__ = [
     "all",
     "ptp",
     "power",
+    "logsumexp",
 ]

tests/tensor/test_math.py

@@ -9,6 +9,7 @@ from itertools import product
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
+from scipy.special import logsumexp as scipy_logsumexp
 
 import aesara.scalar as aes
 from aesara.compile.debugmode import DebugMode
@@ -75,6 +76,7 @@ from aesara.tensor.math import (
     log1p,
     log2,
     log10,
+    logsumexp,
     max,
     max_and_argmax,
     maximum,
@@ -3273,3 +3275,35 @@ def test_tanh_grad_broadcast():
     grad(tanh(x).sum(), x)
     grad(tanh(x + y).sum(), y)
     grad(tanh(x + y).sum(), [x, y])
+
+
+@pytest.mark.parametrize(
+    ["shape", "axis"],
+    [
+        ((1,), 0),
+        ((3,), 0),
+        ((3, 4), None),
+        ((3, 4), 0),
+        ((3, 4), 1),
+        ((3, 4, 5), None),
+        ((3, 3, 5), 0),
+        ((3, 4, 5), 1),
+        ((3, 4, 5), 2),
+    ],
+)
+@pytest.mark.parametrize(
+    "keepdims",
+    [True, False],
+)
+def test_logsumexp(shape, axis, keepdims):
+    scipy_inp = np.zeros(shape)
+    scipy_out = scipy_logsumexp(scipy_inp, axis=axis, keepdims=keepdims)
+
+    aesara_inp = as_tensor_variable(scipy_inp)
+    f = function([], logsumexp(aesara_inp, axis=axis, keepdims=keepdims))
+    aesara_out = f()
+
+    np.testing.assert_array_almost_equal(
+        aesara_out,
+        scipy_out,
+    )