How does PyTorch Tensor.index_select() evaluates tensor output?
I am not able to understand how complex indexing - non contiguous indexing of a tensor works. Here is a sample code and its output
import torch

def describe(x):
  print(&quot;Type: {}&quot;.format(x.type()))
  print(&quot;Shape/size: {}&quot;.format(x.shape))
  print(&quot;Values: \n{}&quot;.format(x))


indices = torch.LongTensor([0,2])
x = torch.arange(6).view(2,3)
describe(torch.index_select(x, dim=1, index=indices))

Returns output as

Type: torch.LongTensor Shape/size: torch.Size([2, 2]) Values:
tensor([[0, 2],
[3, 5]])

Can someone explain how did it arrive to this output tensor?
Thanks!

You are selecting the first (indices[0] is 0) and third (indices[1] is 2) tensors from x on the first axis (dim=0). Essentially, torch.index_select with dim=1 works the same as doing a direct indexing on the second axis with x[:, indices].
&gt;&gt;&gt; x
tensor([[0, 1, 2],
        [3, 4, 5]])

So selecting columns (since you're looking at dim=1 and not dim=0) which indices are in indices. Imagine having a simple list [0, 2] as indices:
&gt;&gt;&gt; indices = [0, 2]

&gt;&gt;&gt; x[:, indices[0]] # same as x[:, 0]
tensor([0, 3])

&gt;&gt;&gt; x[:, indices[1]] # same as x[:, 2]
tensor([2, 5])

So passing the indices as a torch.Tensor allows you to index on all elements of indices directly, i.e. columns 0 and 2. Similar to how NumPy's indexing works.
&gt;&gt;&gt; x[:, indices]
tensor([[0, 2],
        [3, 5]])


Here's another example to help you see how it works. With x defined as x = torch.arange(9).view(3, 3) so we have 3 rows (a.k.a. dim=0) and 3 columns (a.k.a. dim=1).
&gt;&gt;&gt; indices
tensor([0, 2]) # namely 'first' and 'third'

&gt;&gt;&gt; x = torch.arange(9).view(3, 3)
tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])

&gt;&gt;&gt; x.index_select(0, indices) # select first and third rows
tensor([[0, 1, 2],
        [6, 7, 8]])

&gt;&gt;&gt; x.index_select(1, indices) # select first and third columns
tensor([[0, 2],
        [3, 5],
        [6, 8]])

Note: torch.index_select(x, dim, indices) is equivalent to x.index_select(dim, indices)
