Sum at each position on boolean mask without for loop

1

Hi there,

This is my first question, I hope that I’m in the right category!

Briefly, I have a 2d boolean mask, and I want to obtain the sum at each position.

Example:

mask = torch.Tensor([[1, 0, 0, 1, 1, 0, 1],[0, 1, 0, 1, 0, 1, 0]])

A normal sum on axis 1 will be:

torch.sum(mask, dim=1)
tensor([4., 3.])

But I would like to obtain the sum at each position, where the result has the same shape as mask:

sum_pos(mask, dim=1)
tensor([[1, 1, 1, 2, 3, 3, 4],[0, 1, 1, 2, 2, 3, 3]])

Is there any way to do that, without a for loop on dim=1 ?

Thanks in advance!

2

You could use tensor.cumsum() to achieve this.

mask = torch.Tensor([[1, 0, 0, 1, 1, 0, 1],[0, 1, 0, 1, 0, 1, 0]])
mask.cumsum(dim=1)
1 Like