Retrieve element from a vector with an index matrix

Hello,

I am wondering what will be the most efficient way to do the following task in Pytorch?

I have a 2-d matrix consists of the index location, I want to use this matrix to do an element-wise indexing of a vector to get a new 2-d matrix that now contains the actual value.

As an example

2d_index = [[0,3] ] [2,1]]
num_vector = [0, 5, 10, 20]

final_number should be [[0,20][10,5]]

Thanks for the help!

index = torch.tensor([[0,3],[2,1]])
array1d = torch.tensor([0, 5, 10, 20])
selection = array1d[index]
2 Likes

Thanks! This works much faster and is much cleaner than what I have right now. This was what I had before:

index = torch.tensor([[0,3],[2,1]])
array1d = torch.tensor([0, 5, 10, 20])
selection = array1d.reshape(-1).repeat(2,1).gather(1,index)

Yours is ~ 2.8X faster on my CPU

Thanks for this additional info. I have not timed it anytime. It’s good to know.