masked_scatter but rowwise?

user_240443 0
Viewed 473

Assuming a mask as follows:

mask = torch.tensor([
  [True,  True,  False, True,  False],
  [True,  False, True,  True,  True ],
])

I would like to number the True values with sequential values in each row separately. I don't care what's in the False spots, so 0 for simplicity. Thus the desired result is

tensor([[0, 1, 0, 2, 0],    # 0 1 _ 2 _
        [0, 0, 1, 2, 3]])   # 0 _ 1 2 3

I hoped this would work:

replacements = torch.arange(mask.size(1)).expand(mask.size())
target = torch.zeros(mask.size(), dtype=int)
target.masked_scatter(mask, replacements)

Unfortunately, masked_scatter ignores the shape of replacements, so this code results in:

tensor([[0, 1, 0, 2, 0],    # 0 1 _ 2 _
        [3, 0, 4, 0, 1]])   # 3 _ 4 0 1

What would I need to do instead?

pytorch python
1 Answers

I would try something with torch.cumsum: torch.cumsum(mask,dim=1) -1) * mask

The complete example

import torch
mask = torch.tensor([
  [True,  True,  False, True,  False],
  [True,  False, True,  True,  True ],
])
result=torch.cumsum(mask,dim=1) -1) * mask
print(result)

That would print:

tensor([[0, 1, 0, 2, 0],
        [0, 0, 1, 2, 3]])
user_14785370
Related