🚀 Feature

Add numerically stable cumulative logsumexp function. Also we have associated PR on cummax that is needed for numerically stable implementation (#20240).

Motivation

This is useful when computing sum of probabilities and have different applications.

Pitch

Torch has cumsum and cumprod so I suggest logcumsumexp to be added.

Alternatives

Current workaround for me and for people that need it can be written as follows (it is quite fast, only overhead for converting between numpy/torch):

import numpy as np
import torch

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

UPDATE: code above is not always numerically stable. Still most numerically stable, but slow way to compute logcumsumexp is use for-loop:

def logcumsumexp(x, dim):
    # slow implementation, but ok for now
    if (dim != -1) or (dim != x.ndimension() - 1):
        x = x.transpose(dim, -1)

    out = []
    for i in range(1, x.size(-1) + 1):
        out.append(torch.logsumexp(x[..., :i], dim=-1, keepdim=True))
    out = torch.cat(out, dim=-1)

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