Pytorch - Indexing a range of multiple Indices?

user_1436508 0
Viewed 2488

Lets say I have a tensor of size [100, 100] and I have a set of start_indices and end_indices of size [100]

I want to be able to do something like this:

tensor[start_indices:end_indices, :] = 0

Unfortunately, I get an error saying

TypeError: only integer tensors of a single element can be converted to an index

So is this actually possible without a for loop?

pytorch
2 Answers

To the best of my knowledge this is not possible without some sort of loop or list comprehension.

Below are some alternatives which may be useful depending on your use-case. Specifically if you are looking to reuse the same start_indices and end_indices for multiple assignments, or if you are looking have only one in-place assignment to tensor then the solutions below would be useful.

If instead of start_indices and end_indices you were given a list of indices, for example

row_indices = torch.cat([torch.arange(s, e, dtype=torch.int64) for s, e in zip(start_indices, end_indices)])

Then this would be possible using

tensor[row_indices, :] = 0

Or if you were given a mask

mask = torch.zeros(tensor.shape, dtype=torch.bool, device=tensor.device)
for s, e in zip(start_indices, end_indices):
    mask[s:e, :] = True

then this would be possible using

tensor[mask] = 0
user_27900470

To expand a bit on @jodag's answer:

Using torch.cat to create the combined range first (even if they overlap) seems to be some 30% faster than a simple for loop. My tests indicate this holds over a wide range of values for n_ranges (I tried from 100 to 10,000):

tensor_size = 100
n_ranges = 10000
start_indices = [random.randint(0, tensor_size // 2) for _ in range(n_ranges)]
end_indices = [start_indices[i] + random.randint(1, tensor_size // 2) for i in range(n_ranges)]

# 32 milliseconds for this line
row_indices = torch.cat([torch.arange(s, e, dtype=torch.int64)
                        for s, e in zip(start_indices, end_indices)])
# and 3 milliseconds for this line
tensor[row_indices, :] = 0

# 45 milliseconds for this loop; an increase of 28%
for s, e in zip(start_indices, end_indices):
    tensor[s:e, :] = 0

The mask method, surprisingly, is a further 11% time increase compared to the simple for loop when n_ranges = 100 but is 5% faster at n_ranges = 10000.

user_1495060
Related