Change values in specific tensor positions

1

Hello,

I have 2D tensors, where each row starts with 1 and continues with a variable length of 1s followed by 0s.
I want to turn the first and the last 1 of each row to 0.
For instance, for a tensor with torch.Size([3, 7]), that initially is

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

and I want to transform it to the tensor

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

What is the most straight forward way to make this transformation?

2

This should work:

x = torch.tensor([[1, 1, 1, 1, 0, 0, 0],
                  [1, 1, 1, 0, 0, 0, 0],
                  [1, 1, 1, 1, 1, 0, 0]])

tmp = (x == 1).cumsum(dim=1) - 1
idx_start = tmp[:, 0] 
idx_stop = tmp[:, -1]
idx = torch.stack((idx_start, idx_stop), dim=1)
x[torch.arange(x.size(0)).unsqueeze(1), idx] = 0

print(x)
> tensor([[0, 1, 1, 0, 0, 0, 0],
          [0, 1, 0, 0, 0, 0, 0],
          [0, 1, 1, 1, 0, 0, 0]])
2 Likes
3

Thank you @ptrblck, this is a straight forward solution, since you build the tensor positions that are about to change.

That was my first attempt too, and while trying, I thought that maybe there is a build-in function doing that. Anyway, it is a nice solution.