Reverse Cumulative Sum

[lematt1991]

🚀 Feature

Add reverse option to torch.cumsum, such as in tensorflow

Motivation

This would compute right to left cumulative sum more efficiently. Currently, as far as I know, the only way to do it is

x = torch.arange(9).view(3, 3)
r2lcumsum = th.flip(th.cumsum(th.flip(x, [1]), 1), [1])

Result should be:

tensor([[ 3,  3,  2],
        [12,  9,  5],
        [21, 15,  8]])

Pitch

Add reverse arg to native cumsum function

Alternatives

Additional context

[vadimkantorov]

also stumbled on this. @lematt1991 thanks for the workaround!

[xuuuluuu]

also add reverse for cumprod function

[yufeng66]

Yes. Please add this. Using the workaround doubles the run time

[yash-s20]

Here's a faster workaround

import torch
x = torch.arange(9).view(3, 3)
r2lcumsum = x +  torch.sum(x, dim=1, keepdims=True) - torch.cumsum(x, dim=1)

This completely avoids the need for flip, which seems to be the bottleneck
Hope this helps someone!

Didn't check but should also work for cumprod, simply replace sum with prod, addition with multiplication, subtraction with division (unless of course dealing with zeros)

[Yonv1943]

Base on the answer of @yash-s20

We don't have to calculate torch.sum(x) again, because torch.cumsum(x)[-1] == torch.sum(x).

So we have:

import torch

x = torch.arange(1, 4)
print('[1, 2, 3]', x)

cum_sum = torch.cumsum(x, dim=0)             # faster way
re_cum_sum = x - cum_sum + cum_sum[-1:None]  # faster way
print('[1, 3, 6]', cum_sum)
print('[6, 5, 3]', re_cum_sum)

re_cum_sum = torch.cumsum(x.flip(dims=(0,)), dim=0).flip(dims=(0,))  # slower way
print('[6, 5, 3]', re_cum_sum)

[rtqichen]

Small note here that the ``faster'' way of subtracting cumsum from the full sum can introduce additional roundoff errors because it relies on adding and subtracting many unwanted terms. This has consequences if what you're doing requires high numerical precision. It would still be ideal to have a native implementation of revcumsum as it would be both faster and more stable.