python - Index multidimensional torch tensor by another multidimensional tensor

3
votes

I have a tensor x in pytorch let's say of shape (5,3,2,6) and another tensor idx of shape (5,3,2,1) which contain indices for every element in first tensor. I want a slicing of the first tensor with the indices of the second tensor. I tried x= x[idx] but I get a weird dimensionality when I really want it to be of shape (5,3,2) or (5,3,2,1).

I'll try to give an easier example: Let's say

x=torch.Tensor([[10,20,30],
                 [8,4,43]])
idx = torch.Tensor([[0],
                    [2]])

I want something like

y = x[idx]

such that 'y' outputs [[10],[43]] or something like.

The indices represent the position of the wanted elements the last dimension. for the example above where x.shape = (2,3) the last dimension are the columns, then the indices in 'idx' is the column. I want this but for more than 2 dimensions

2
pythonnumpypytorchtensornumpy-slicing
How do interpret indices idx=[[0],[2]] to get values [[10],[43]] from x? It is unclear what those indices represent, are they row/column or flattened array indices? - Ehsan
It would mean the position in the last dimension which for that example is the column. - Jessica Borja
Please check the post to see if that is what you needed. Thank you. Please also check out stackoverflow.com/help/someone-answers on how to accept answers on SO. - Ehsan

2 Answers

1
votes

From what I understand from the comments, you need idx to be index in the last dimension and each index in idx corresponds to similar index in x (except for the last dimension). In that case (this is the numpy version, you can convert it to torch):

ind = np.indices(idx.shape)
ind[-1] = idx
x[tuple(ind)]

output:

[[10]
 [43]]
0
votes

You can use range; and squeeze to get proper idx dimension like

x[range(x.size(0)), idx.squeeze()]
tensor([10., 43.])

# or
x[range(x.size(0)), idx.squeeze()].unsqueeze(1)
tensor([[10.],
        [43.]])