Stable and efficient implementation of logcumsumexp

agadetsky: <p>There is an implementation of <code>logsumexp</code> in pytorch, but there is no <code>logcumsumexp</code>. How can it be implemented efficiently and stable in pytorch?</p>

KFrank: <p>Hello Artyom!</p>
<aside class="quote no-group" data-username="agadetsky" data-post="1" data-topic="55886" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/f04885/48.png" class="avatar"> agadetsky:</div>
<blockquote>
<p>There is an implementation of <code>logsumexp</code> in pytorch, but there is no <code>logcumsumexp</code>. How can it be implemented efficiently and stable in pytorch?</p>
</blockquote>
</aside>
<p>As it stands, it cannot, other than someone writing it.</p>
<p>You presumably understand the numerical issues with calculating<br>
<code>logsumexp</code>.  (See, for example, the discussion in Wikipedia’s<br>
<a href="https://en.wikipedia.org/wiki/LogSumExp" rel="noopener nofollow ugc">LogSumExp</a>.)</p>
<p>So if pytorch had, for example, a <code>cummax</code> tensor function, you<br>
could implement <code>logcumsumexp</code> using pytorch tensor functions.</p>
<p>But this doesn’t exist (yet).  See:</p>
<p><a href="https://stackoverflow.com/questions/55665624/vectorized-implementation-of-cumulative-maximum-in-pytorch-with-requires-grad-tr" rel="noopener nofollow ugc">https://stackoverflow.com/questions/55665624/vectorized-implementation-of-cumulative-maximum-in-pytorch-with-requires-grad-tr</a></p>
<p><a href="https://github.com/pytorch/pytorch/issues/20240" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/20240</a></p>
<p>and</p>
<p><a href="https://discuss.pytorch.org/t/sliding-max-over-dimension/49799">https://discuss.pytorch.org/t/sliding-max-over-dimension/49799</a></p>
<p>So, short of writing the <code>logcumsumexp</code> (or related) tensor<br>
function “from scratch,” you would have to use a loop to get<br>
the “running maximum” (<code>cummax</code>) part, thus forgoing some<br>
of the efficiency provided by using just tensor functions.</p>
<p>Good luck.</p>
<p>K. Frank</p>

agadetsky: <p><a class="mention" href="/u/kfrank">@KFrank</a> Thank you for your answer. I created PR, maybe someone will add <code>cummax</code>  and <code>logcumsumexp</code> at ATen to be efficient. At the moment I end up with the following implementation, maybe someone will need it in the future.</p>
<pre><code class="lang-auto">import torch
import numpy as np

def cummax(x, dim):
    x_np = x.detach().cpu().numpy()
    ret = np.maximum.accumulate(x_np, axis=dim)
    return torch.from_numpy(ret).to(x)


def logcumsumexp(x, dim=-1):
    if (dim != -1) or (dim != x.ndimension() - 1):
        x = x.transpose(dim, -1)

    init_size = x.size()
    last_dim_size = init_size[-1]
    x_resized = x.contiguous().view(-1, last_dim_size)
    d1, d2 = x_resized.size()
    x_cummax = cummax(x_resized, -1).view(d1, d2, 1)
    x_expand = x_resized.unsqueeze(1).expand(d1, d2, last_dim_size)
    mask = torch.tril(torch.ones(last_dim_size, last_dim_size)).unsqueeze(0)
    ret = torch.log(torch.sum(torch.exp(x_expand - x_cummax) * mask, dim=-1)) + x_cummax.view(d1, d2)
    ret = ret.view(*init_size)

    if (dim != -1) or (dim != x.ndimension() - 1):
        ret = ret.transpose(-1, dim)

    return ret
</code></pre>