torch.select(): Slicing a Tensor in PyTorch

The torch.select() method in PyTorch slices the input tensor along a specified dimension at the given index. It retrieves a single slice of a tensor, reducing the dimensionality of the output by one compared to the input tensor.

The main use case of this method is when you want to narrow a tensor along a specific dimension. 

The output tensor is a view of the input tensor (no data copy), and has one less dimension than the input. For example, a 3D tensor (2×3×4) becomes 2D (2×4) after selecting along dim=1.

Syntax

torch.select(input, dim, index)

Parameters

Arguments	Description
input (Tensor)	It represents an input tensor from which you need to select the slice.
dim (Tensor)	It is the dimension along which to choose the slice. It can be positive (0-based indexing) or negative (counting from the end).
index (Tensor)	It is an index of the slice you can select along the specified dimension.

Selecting a slice from a 1D tensor

In a 1D tensor, there is only one dimension, so the dim argument for torch.select() must be 0

To select element 21, we passed the index = 2, which is the third element of the tensor.

Let’s implement the above figure into a program.

import torch

tensor = torch.tensor([11, 21, 19, 10])

selected = tensor.select(dim=0, index=2)

print(selected)

# Output: tensor(19)

At index 2, the element is 19, so it selects 19.

Selecting a column in a 2D tensor (Matrix)

 

To select a column from a 2D tensor, pass dim=1. It refers to the column dimension.

 For specific selection, you also need to pass its index. Remember that the column index starts with 0.

import torch

tensor_2d = torch.tensor([[12, 2, 31],
                          [41, 5, 62]])

column_selection = torch.select(tensor_2d, dim=1, index=1)

print(column_selection)

# Output: tensor([2, 5])

Let me explain the above code:

1. The first column values are 12 and 41, and the index is 0.
2. The second column values are 2 and 5, and the index is 1.
3. The third column values are 31 and 62, and the index is 2.

We selected the second column, so the output tensor values are [2, 5].

Selecting a row from a 2D Tensor (Matrix)

To select a row, you need to pass dim=0. It refers to a row dimension.

import torch

tensor_2d = torch.tensor([[12, 2, 31],
                          [41, 5, 62]])

row_selection = torch.select(tensor_2d, dim=0, index=1)

print(row_selection)

# Output: tensor([41,  5, 62])

Let me explain the above code:

1. The row index starts with 0 and has 12, 2, and 31 values. 
2. The second row begins with 1 and its values are 41, 5, and 62.

Since we pass the index 1, we select the second row in the output.

Using negative dimension

Negative dimension refers to the last dimension (columns in a 2D tensor).

If you pass -1 in a 2D tensor, it will return the tensor column-wise.

import torch

tensor_2d = torch.tensor([[12, 2, 31],
                          [41, 5, 62]])

negative_column = torch.select(tensor_2d, dim=-1, index=2)

print(negative_column)

# Output: tensor([31, 62])

Index 2 refers to the third column, and index -1 refers to the row. Hence, the third column has been selected.

import torch

tensor_2d = torch.tensor([[12, 2, 31],
                          [41, 5, 62]])

out_of_bound_index = torch.select(tensor_2d, dim=0, index=2)

print(out_of_bound_index)

# Output: 

# IndexError: select(): index 2 out of range for tensor of size [2, 3] at dimension 0

Ensure that you are passing a valid dimension. Here, we checked for the index.

What about dimension? What if we pass the dimension that does not exist?

import torch

tensor_2d = torch.tensor([[12, 2, 31],
                          [41, 5, 62]])

out_of_bound_dimension = torch.select(tensor_2d, dim=2, index=1)

print(out_of_bound_dimension)

# Output:
# IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)

That’s it!