Change tensor values by index greater(or less) than some value

1

Hey everyone!

I have a question of selecting/changing values by some specific indices.
Here is an example:

import torch

"""
The original tensor is like this:

t = torch.tensor(
    [[99, 2, 1, 3, 4],
     [1, 2, 3, 99, 4],
     [4, 1, 99, 2, 3]]
)

I want to change t's value with rule:
for each row, set elements which is before 99 to 1; 0 otherwise,

so my target is:

tensor(
    [[0, 0, 0, 0],
    [1, 1, 1, 0], 
    [1, 1, 0, 0]]
)

following is the for-loop way, my question is, can I do this more efficiently?

Thanks!

"""

# create a tensor of shape [3, 4]
t = torch.tensor(
    [[99, 2, 1, 3, 4],
     [1, 2, 3, 99, 4],
     [4, 1, 99, 2, 3]]
)

# get 99's position of each row
position_of_zero = torch.nonzero(t == 99)[:, 1]

print(position_of_zero)

# for-loop
for row, zero_pos in zip(t, position_of_zero):
    for i in range(zero_pos):
        row[i] = 1
    
    for i in range(zero_pos + 1, len(row)):
        row[i] = 0

print(t)

# drop 99
t = t[t != 99].view(-1, 4)
print(t)
2

The desired target doesn’t match the description:

for each row, set elements which is before 99 to 1; 0 otherwise,

since the 99 in the first row is set to 0, while the others are not.
Assuming you would like to set the value at 99 to 1, this code should work and would avoid the loop:

t2 = torch.tensor(
    [[99, 2, 1, 3, 4],
     [1, 2, 3, 99, 4],
     [4, 1, 99, 2, 3]]
)

idx = torch.cumsum((t2 == 99), 1)
idx[t2==99] = 0
res = (~idx.bool()).long()
print(res)
> tensor([[1, 0, 0, 0, 0],
          [1, 1, 1, 1, 0],
          [1, 1, 1, 0, 0]])

Depending on the shape of your tensor, it might be faster or not, so I would recommend to profile both approaches.

3

since the 99 in the first row is set to 0 , while the others are not.

Sorry for the unclear description, actually I want to just DROP all the 99s,

I slightly modified you solution, and it works for me too!

t2 = torch.tensor(
    [[99, 2, 1, 3, 4],
     [1, 2, 3, 99, 4],
     [4, 1, 99, 2, 3]]
)

idx = torch.cumsum((t2 == 99), 1)
idx[t2==99] = 2
idx = idx[idx != 2].view(3, -1)

res = (~idx.bool()).long()
print(res)
> tensor([[1, 0, 0, 0, 0],
          [1, 1, 1, 1, 0],
          [1, 1, 1, 0, 0]])

Thanks for your help!