How does `torch.gather()` construct new tensors?
I'm struggling to understand what the function torch.gather is doing.
If I have:
a = torch.tensor([[4,5,6],[7,8,9],[10,11,12]])
b = torch.tensor([[1,1],[1,2]])

Then
&gt;&gt;&gt; torch.gather(input=a, dim=1, index=b)
tensor([[5, 5],
        [8, 9]])

While
&gt;&gt;&gt; torch.gather(input=a, dim=0, index=b)
tensor([[ 7,  8],
        [ 7, 11]])

Can somebody explain very simply how these output tensors are actually constructed?

On the first call you index a with b along dim=1 (2nd dimension). The performed operation is:
out[i,j] = a[i, b[i,j]]

Which returns:
[[ a[0, b[0,0]], a[0, b[0,1]] ],
 [ a[1, b[1,0]], a[1, b[1,1]] ]]

While on the second call, you index a with b along dim=0 (1st dimension). This time around:
out[i,j] = a[b[i,j], j]

Which returns:
[[ a[b[0,0], 0], a[b[0,1], 1] ],
 [ a[b[1,0], 0], a[b[1,1], 1] ]]

For more explanations, read on with this thread.
