Select items from a matrix or tensor using 1-D array of indexes to get 1D or 2D tensors in torch

I have a tensor of n-sampled predictions. In this example I sample 4 times a (10,2) result, where this results represents a batch of graphs with each graph having 3+4+3=10 nodes with 2 coordinates. I have a way to select the best sample for each graph, and I have an index telling me this (idx_test). I want to select in an efficient way the elements of the sampled tensor. The best way I have found is like this (toy code example):

rand_test = torch.randint(3, 10, (4, 10, 2))
idx_test = torch.LongTensor([0, 3, 1])

n_nodes = [3, 4, 3]

final_tensor = torch.zeros((10, 2))
accum_nodes = torch.cumsum(torch.LongTensor(n_nodes), 0)
# First iteration is done outside the for loop
final_tensor[:accum_nodes[0]] = rand_test[idx_test[0],:accum_nodes[0],:]
for i, idx in enumerate(idx_test):
    final_tensor[accum_nodes[i-1]:accum_nodes[i]] = rand_test[idx,accum_nodes[i-1] :accum_nodes[i], :]
print(final_tensor)

The objective is to obtain the same final tensor without using for loops since in my model this array is large

Thanks!

Answer

You could create a helper tensor containing the ranges of node indices you want to index rand_test with. In your example, you have: i=0 you have 3 nodes with index value 0, i=1 has 4 with value 3, while i=2 has 3 with index value 1.

You can do so using repeat_interleave:

>>> idx_test.repeat_interleave(torch.tensor(n_nodes))
tensor([0, 0, 0, 3, 3, 3, 3, 1, 1, 1])

Then index rand_test using idx and a range:

>>> final_tensor = rand_test[idx, range(len(idx))]
tensor([[5, 5],
        [7, 3],
        [8, 4],
        [7, 5],
        [7, 6],
        [7, 8],
        [9, 9],
        [7, 7],
        [8, 7],
        [3, 7]])