How to create a mask tensor from another mask tensor

1

lets say I have tensor of shape (2, 10) .
I want to create a mask tensor with same shape that is ‘1’ only immediately after occurrence of a specific value and ‘0’ elsewhere.
For example considering I want to mask indices after 3:

x = tensor([[1, 3, 4, 5, 3, 6],
        [1, 2, 5, 3, 7, 9]])
mask = tensor([[False,  False, True, False,  False, True],
        [False, False, False,  False, True, False]])

x.eq(3) will mask indices of 3 but not the next index.

Thanks for ur help

2

Just x== value, you can get the result you want

3

I want the next element of the index with that specific value be ‘True’ not the index of value

4

You can return the index of value by torch.nonzero(x==3), then manipulate the index, to generate the mask.

5

Yeah actually I tried that but I could not generate what I want since each row may have multiple occurrence of that value. I don’t know how to manipulate the result of torch.nonzero(x==3).

ind = (x==3).nonzero()
tensor([[0, 1],
        [0, 4],
        [1, 3]])

Should I do it in a for loop? Or is there a nice vectorized way of doing it?

6

===== 1

mask = x == 3
mask = torch.cat([torch.zeros(mask.shape[0])[..., None].bool(), mask[:, :-1]], dim=1)

===== 2

row, col = torch.where(x == 3)
col += 1
mask = torch.zeros_like(x).bool()
mask[row, col] = True

===== 3

mask = torch.zeros_like(x).bool()
mask[:, 1:] = (x[..., :-1] == 3)
1 Like
7

Great, thanks that worked.