Indexing using pyTorch tensors along one specific dimension with 3 dimensional tensor
I have 2 tensors:
A with shape (batch, sequence, vocab)
and B with shape (batch, sequence).
A = torch.tensor([[[ 1.,  2.,  3.],
     [ 5.,  6.,  7.]],

    [[ 9., 10., 11.],
     [13., 14., 15.]]])

B = torch.tensor([[0, 2],
    [1, 0]])

I want to get the following:
C = torch.zeros_like(B)
for i in range(B.shape[0]):
   for j in range(B.shape[1]):
      C[i,j] = A[i,j,B[i,j]]

But in a vectorized way. I tried torch.gather and other stuff but I cannot make it work.
Can anyone please help me?

&gt;&gt;&gt; import torch
&gt;&gt;&gt; A = torch.tensor([[[ 1.,  2.,  3.],
...      [ 5.,  6.,  7.]],
... 
...     [[ 9., 10., 11.],
...      [13., 14., 15.]]])
&gt;&gt;&gt; B = torch.tensor([[0, 2],
...     [1, 0]])
&gt;&gt;&gt; A.shape
torch.Size([2, 2, 3])
&gt;&gt;&gt; B.shape
torch.Size([2, 2])
&gt;&gt;&gt; C = torch.zeros_like(B)
&gt;&gt;&gt; for i in range(B.shape[0]):
...    for j in range(B.shape[1]):
...       C[i,j] = A[i,j,B[i,j]]
... 
&gt;&gt;&gt; C
tensor([[ 1,  7],
        [10, 13]])
&gt;&gt;&gt; torch.gather(A, -1, B.unsqueeze(-1))
tensor([[[ 1.],
         [ 7.]],

        [[10.],
         [13.]]])
&gt;&gt;&gt; torch.gather(A, -1, B.unsqueeze(-1)).shape
torch.Size([2, 2, 1])
&gt;&gt;&gt; torch.gather(A, -1, B.unsqueeze(-1)).squeeze(-1)
tensor([[ 1.,  7.],
        [10., 13.]])

Hi, you can use torch.gather(A, -1, B.unsqueeze(-1)).squeeze(-1).
the first -1 between A and B.unsqueeze(-1) is indicating the dimension along which you want to pick the element.
the second -1 in B.unsqueeze(-1) is to add one dim to B to make the two tensor the same dims otherwise you get RuntimeError: Index tensor must have the same number of dimensions as input tensor.
the last -1 is to reshape the result from torch.Size([2, 2, 1]) to torch.Size([2, 2])
