Set values of 1D tensor using list of start/end indices

#1

Given a 1D tensor x, and a 2D tensor (N * 2) of start and end indices, is there a way to set the values in the start and end range to a specific value (say 0)?
I’m looking for a vectorized way to do the following:

>>> x = torch.randint(10, (8,))
>>> x
tensor([1, 6, 1, 9, 3, 2, 6, 2])
>>> indices = torch.tensor([[1,2], [4,6]])
>>> indices
tensor([[1, 2],
[4, 6]])
>>> desired_output = torch.tensor([1, 0, 1, 9, 0, 0, 6, 2])
>>> desired_output
tensor([1, 0, 1, 9, 0, 0, 6, 2])

The range between start and end is not guaranteed to have the same length (as 1 and 2 here). Any help would be appreciated!

#2

You could create a mask using cumsum and set the values to zero as seen here:

x = torch.tensor([1, 6, 1, 9, 3, 2, 6, 2])

indices = torch.tensor([[1,2], [4,6]])
val = torch.tensor([[1, -1], [1, -1]])
mask = torch.zeros_like(x).scatter_(
    0, indices.view(-1), val.view(-1)).cumsum(0).bool()

x[mask] = 0
print(x)
> tensor([1, 0, 1, 9, 0, 0, 6, 2])

Depending on the needed flexibility of the input tensors you could simplify the val creation etc.

1 Like
#3

Thanks, this is brilliant!

#4

@ptrblck Sorry I have a follow up question: the solution seems to assume that the start and end ranges won’t overlap with each other. Is there a way to modify it to work for overlapping cases? For example, when indices = torch.tensor([[1,3], [2,3]]) or indices = torch.tensor([[1,3], [1,2]]).