PyTorch Tensor.index_select() 如何评估张量 output？

Question

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("Type: {}".format(x.type()))
  print("Shape/size: {}".format(x.shape))
  print("Values: \n{}".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!

Answer 1

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] .

>>> 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 :

>>> indices = [0, 2]

>>> x[:, indices[0]] # same as x[:, 0]
tensor([0, 3])

>>> 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, ie columns 0 and 2 . Similar to how NumPy's indexing works.

>>> 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 (aka dim=0 ) and 3 columns (aka dim=1 ).

>>> indices
tensor([0, 2]) # namely 'first' and 'third'

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

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

>>> 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)