Repository: pytorch/pytorch

Issue:
Feature request: logsumexp
The numerically stable version of `logsumexp` is simple but often useful. Thoughts on including this in PyTorch?

My thought is that it'd follow a call signature(s) that's similar to that of `sum`, but `logsumexp` wouldn't have the `out` keyword, and it'd also use the NumPy-consistent `keepdims` keyword instead of `sum`'s `keepdim` keyword. So

```
def logsumexp(inputs, dim=None, keepdims=False):
    ...
```


Discussion:
Just a couple of questions, might be extremely naive:
- Is this supposed to be differentiable? As in, if it is only going to be used in the `forward()` pass, then is a library function needed? I cannot think of a scenario where I might want to call `backward()` on this, so I'm asking if that is a concern.
- Why `keepdims` instead of `keepdim`? Let's keep the notation consistent across PyTorch!
Yes it should be differentiable :). Many use cases exist, one being mixture density networks.

I'm all for consistency. My best guess is that PyTorch accidentally used keepdim, but that they will strive for numpy consistency in the future. Same happened with TensorFlow. So my thought is that any new functionality should meet this goal in advance.
We used keepdim because reduction functions in pytorch only support a single dimension.  There are plans/patches to change that, after which keepdims would be more appropriate.
Assuming that `log_softmax` is numerically stable, here is a one-line workaround:
```python
import torch.nn.functional as F

def logsumexp(inputs, dim=None, keepdim=False):
    return (inputs - F.log_softmax(inputs)).mean(dim, keepdim=keepdim)
```
Perhaps not very efficient for large sums, but it delegates everything to `log_softmax` as well as replicates its behavior (#1020). Having a native implementation down the line would be nice.
Here is an implementation that we wrote, which supports both variable and tensor inputs:
```python
def log_sum_exp(value, dim=None, keepdim=False):
    """Numerically stable implementation of the operation

    value.exp().sum(dim, keepdim).log()
    """
    # TODO: torch.max(value, dim=None) threw an error at time of writing
    if dim is not None:
        m, _ = torch.max(value, dim=dim, keepdim=True)
        value0 = value - m
        if keepdim is False:
            m = m.squeeze(dim)
        return m + torch.log(torch.sum(torch.exp(value0),
                                       dim=dim, keepdim=keepdim))
    else:
        m = torch.max(value)
        sum_exp = torch.sum(torch.exp(value - m))
        if isinstance(sum_exp, Number):
            return m + math.log(sum_exp)
        else:
            return m + torch.log(sum_exp)
```
I agree that an implementation in C would be very helpful! 
I'm looking for this function as well. I'm implementing the CTC loss, and need to add lots of log probabilities.
+1.  Additionally I have a situation where I need row-wise logsumexp on a sparse matrix (treating unstored entries as if they are -inf rather than zero).  I can't find any way to construct this to be done efficiently on the GPU.
Throwing my version here too. Not sure if there are any tradeoffs with the version above, though the name is consistent with scipy's `logsumexp`, it's a bit shorter, and `dim=None` works.

```
def logsumexp(inputs, dim=None, keepdim=False):
    """Numerically stable logsumexp.

    Args:
        inputs: A Variable with any shape.
        dim: An integer.
        keepdim: A boolean.

    Returns:
        Equivalent of log(sum(exp(inputs), dim=dim, keepdim=keepdim)).
    """
    # For a 1-D array x (any array along a single dimension),
    # log sum exp(x) = s + log sum exp(x - s)
    # with s = max(x) being a common choice.
    if dim is None:
        inputs = inputs.view(-1)
        dim = 0
    s, _ = torch.max(inputs, dim=dim, keepdim=True)
    outputs = s + (inputs - s).exp().sum(dim=dim, keepdim=True).log()
    if not keepdim:
        outputs = outputs.squeeze(dim)
    return outputs
```
Yet another version that handles `inf` and `-inf` cases properly so that no `nan` appears. Requires `torch.where` from nightly build or work-around for version 0.3:
```python
# https://github.com/pytorch/pytorch/issues/2591
import torch

def logsumexp(x, dim=None, keepdim=False):
    if dim is None:
        x, dim = x.view(-1), 0
    xm, _ = torch.max(x, dim, keepdim=True)
    x = torch.where(
        (xm == float('inf')) | (xm == float('-inf')), 
        xm,
        xm + torch.log(torch.sum(torch.exp(x - xm), dim, keepdim=True)))
    return x if keepdim else x.squeeze(dim)
```
A possible `torch.where` work-around for version 0.3 which does not handle broadcasting:
```python
def where(cond, xt, xf):
    ret = torch.zeros_like(xt)
    ret[cond] = xt[cond]
    ret[cond ^ 1] = xf[cond ^ 1]
    return ret
```


Resolving patch:
diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py
--- a/torch/_tensor_docs.py
+++ b/torch/_tensor_docs.py
@@ -1143,6 +1143,13 @@ def add_docstr_all(method, docstr):
     f(x) = \\dfrac{1}{x \\sigma \\sqrt{2\\pi}}\ e^{-\\dfrac{(\\ln x - \\mu)^2}{2\\sigma^2}}
 """)
 
+add_docstr_all('logsumexp',
+               r"""
+logsumexp(dim, keepdim=False) -> Tensor
+
+See :func:`torch.logsumexp`
+""")
+
 add_docstr_all('lt',
                r"""
 lt(other) -> Tensor
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -2154,6 +2154,36 @@ def parse_kwargs(desc):
     tensor([  1.2589,   2.1135,   3.5481,   5.9566,  10.0000])
 """.format(**factory_common_args))
 
+add_docstr(torch.logsumexp,
+           r"""
+logsumexp(input, dim, keepdim=False, out=None)
+
+Returns the log of summed exponentials of each row of the :attr:`input`
+tensor in the given dimension :attr:`dim`. The computation is numerically
+stabilized.
+
+For summation index :math:`j` given by `dim` and other indices :math:`i`, the result is
+
+           :math:`\text{logsumexp}(x)_{i} = \log \sum_j \exp(x_ij).`
+
+If :attr:`keepdim` is ``True``, the output tensor is of the same size
+as :attr:`input` except in the dimension :attr:`dim` where it is of size 1.
+Otherwise, :attr:`dim` is squeezed (see :func:`torch.squeeze`), resulting in
+the output tensor having 1 fewer dimension than :attr:`input`.
+
+Args:
+    input (Tensor): the input tensor
+    dim (int or tuple of ints): the dimension or dimensions to reduce
+    keepdim (bool): whether the output tensor has :attr:`dim` retained or not
+    out (Tensor, optional): the output tensor
+
+
+Example::
+    >>> a = torch.randn(3, 3)
+    >>> torch.logsumexp(a, 1)
+    tensor([ 0.8442,  1.4322,  0.8711])
+""")
+
 add_docstr(torch.lt,
            r"""
 lt(input, other, out=None) -> Tensor


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -2738,6 +2738,7 @@ class dont_convert(tuple):
     ('addcdiv', (), (0.5, (S, S, 1), (1, S)), 'scalar_scale_broadcast_lhs'),
     ('zero_', (S, S, S), NO_ARGS),
     ('zero_', (), NO_ARGS, 'scalar'),
+    ('logsumexp', (S, S), (1,)),
     ('norm', (S, S), (2,)),
     ('norm', (S, S), (0,), '0'),
     ('norm', (S, S), (0.5,), '0_5'),
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -766,6 +766,17 @@ def test_multidim(x, dim):
     def test_dim_reduction(self):
         self._test_dim_reduction(self, lambda t: t)
 
+    @unittest.skipIf(not TEST_SCIPY, "Scipy not found")
+    def test_logsumexp(self):
+        from scipy.special import logsumexp
+        a = torch.randn(5, 4)
+        a[0, 0] = float('inf')
+        a[1, :] = float('-inf')
+        actual = a.logsumexp(1)
+        expected = logsumexp(a.numpy(), 1)
+        self.assertEqual(expected.shape, actual.shape)
+        self.assertTrue(np.allclose(expected, actual.numpy()))
+
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_cpu_parallel(self):
         # To use parallel branches we'll need to compare on tensors