如何实现PyTorch中的多维花式索引？

Question

I have a 3D tensor ( [batch_size, seq_length, hidden_dim] ) and a 2D list ( [batch_size, seq_length] ).

I want to using list to accomplish the selection of this tensor.

For example: the shape of 3D tensor t [2, 5, 3] and the shape of 2D list l [2, 5] .

Let t0 = t[0, :, :] and l0 = l[0] . I would like to select " t0[l0] ".

Same for t1[l1] , and so on.

I can only think of writng a for loop for achieve this:

new_ts = [] 
for i in range(t.shape[0]):
    new_t = t[i][l[i]]
    new_ts.append(new_t) 
new_t = torch.cat(new_ts, dim=2)

There must be a more simple way to accomplish this. I have also tried multi-dimensional fancy indexing t[l] , but the syntax is not valid and it doesn't work.

Looking forward to your suggestions.

Answer 1

The code provided doesn't work, it seems you are looking to concatenate the tensors in the list on an axis that doesn't yet exit. I suggest you use stack instead.

I don't think you can achieve this without a single loop. In terms of syntax and readability you can do a little better, in my opinion...

Having both starting tensors t and l :

t = torch.rand(2,5,3)
l = torch.randint(0,5, size=(2,5))

Using a list comprehension:

torch.stack([t[i][jj] for i, jj in enumerate(l)])