Index-based Operation in PyTorch

Index-based operations play a vital role in manipulating and accessing specific elements or subsets of data within tensors. PyTorch, a popular open-source deep learning framework, provides powerful mechanisms to perform such operations efficiently. By leveraging index-based operations, developers can extract, modify, and rearrange data along various dimensions of a tensor.

Tensor Basics

PyTorch tensors are multi-dimensional arrays that can hold numerical data of various types, such as floating-point numbers, integers, or Boolean values. Tensors are the fundamental data structure in PyTorch and serve as the building blocks for constructing and manipulating neural networks.

To create a tensor in PyTorch, we can use the torch.Tensor class or various factory functions provided by PyTorch, such as torch.zeros, torch.ones, or torch.rand. Let's look at a few examples ?

import torch

# Create a tensor of zeros with shape (3, 2)
zeros_tensor = torch.zeros(3, 2)
print("Zeros tensor:")
print(zeros_tensor)

# Create a tensor of ones with shape (2, 3)
ones_tensor = torch.ones(2, 3)
print("\nOnes tensor:")
print(ones_tensor)

# Create a random tensor with shape (3, 3)
rand_tensor = torch.rand(3, 3)
print("\nRandom tensor:")
print(rand_tensor)

Zeros tensor:
tensor([[0., 0.],
        [0., 0.],
        [0., 0.]])

Ones tensor:
tensor([[1., 1., 1.],
        [1., 1., 1.]])

Random tensor:
tensor([[0.4963, 0.7682, 0.0885],
        [0.1320, 0.3074, 0.6341],
        [0.4901, 0.8964, 0.4556]])

In addition to the tensor's shape, we can also inspect its data type using the dtype attribute. PyTorch supports a wide range of data types, including torch.float32, torch.float64, torch.int32, and torch.bool. To specify a particular data type, we can pass the dtype argument when creating a tensor.

import torch

# Create a tensor with specific data type
ones_double_tensor = torch.ones(2, 2, dtype=torch.float64)
print("Double precision tensor:")
print(ones_double_tensor)
print("Data type:", ones_double_tensor.dtype)

Double precision tensor:
tensor([[1., 1.],
        [1., 1.]], dtype=torch.float64)
Data type: torch.float64

Creating Tensors from Data

We can convert existing data structures, such as lists or NumPy arrays, into PyTorch tensors using the torch.tensor function ?

import torch
import numpy as np

# Create from list
data_list = [[1, 2, 3], [4, 5, 6]]
tensor_from_list = torch.tensor(data_list)
print("Tensor from list:")
print(tensor_from_list)

# Create from NumPy array
numpy_array = np.array([[7, 8, 9], [10, 11, 12]])
tensor_from_numpy = torch.tensor(numpy_array)
print("\nTensor from NumPy:")
print(tensor_from_numpy)

Tensor from list:
tensor([[1, 2, 3],
        [4, 5, 6]])

Tensor from NumPy:
tensor([[ 7,  8,  9],
        [10, 11, 12]])

Basic Indexing and Slicing

Indexing and slicing operations allow us to access specific elements or subsets of tensors. The indexing starts from 0 for the first element in each dimension ?

import torch

# Create a tensor
tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print("Original tensor:")
print(tensor)

# Access individual elements
element = tensor[0, 1]  # Row 0, Column 1
print(f"\nElement at [0, 1]: {element}")

# Access with negative indices
last_element = tensor[-1, -1]  # Last row, last column
print(f"Last element: {last_element}")

# Slice operations
first_row = tensor[0, :]  # All columns of first row
print(f"First row: {first_row}")

first_column = tensor[:, 0]  # All rows of first column
print(f"First column: {first_column}")

# Extract submatrix
submatrix = tensor[1:, 1:]  # From row 1 onwards, column 1 onwards
print(f"Submatrix:\n{submatrix}")

Original tensor:
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])

Element at [0, 1]: 2
Last element: 9
First row: tensor([1, 2, 3])
First column: tensor([1, 4, 7])
Submatrix:
tensor([[5, 6],
        [8, 9]])

Advanced Indexing with Arrays

PyTorch supports advanced indexing using integer arrays to select specific indices ?

import torch

# Create a tensor
tensor = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Use integer array indexing
indices = torch.tensor([0, 2])  # Select rows 0 and 2
selected_rows = tensor[indices]
print("Selected rows (0 and 2):")
print(selected_rows)

# Advanced indexing for specific elements
row_indices = torch.tensor([0, 1, 2])
col_indices = torch.tensor([2, 1, 0])
diagonal_elements = tensor[row_indices, col_indices]
print(f"\nDiagonal elements: {diagonal_elements}")

Selected rows (0 and 2):
tensor([[1, 2, 3],
        [7, 8, 9]])

Diagonal elements: tensor([3, 5, 7])

Boolean Mask Indexing

Boolean masks provide a powerful way to select elements based on conditions ?

import torch

# Create a tensor
tensor = torch.tensor([1, 2, 3, 4, 5, 6])
print(f"Original tensor: {tensor}")

# Create boolean mask
mask = tensor > 3
print(f"Boolean mask (> 3): {mask}")

# Select elements using mask
selected = tensor[mask]
print(f"Selected elements: {selected}")

# Apply to 2D tensor
matrix = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
mask_2d = matrix > 5
print(f"\n2D Boolean mask (> 5):\n{mask_2d}")
print(f"Selected elements from 2D: {matrix[mask_2d]}")

Original tensor: tensor([1, 2, 3, 4, 5, 6])
Boolean mask (> 3): tensor([False, False, False,  True,  True,  True])
Selected elements: tensor([4, 5, 6])

2D Boolean mask (> 5):
tensor([[False, False, False],
        [False, False,  True],
        [ True,  True,  True]])
Selected elements from 2D: tensor([6, 7, 8, 9])

Ellipsis for Extended Slicing

The ellipsis (...) syntax allows convenient slicing of high-dimensional tensors ?

import torch

# Create a 4D tensor
tensor_4d = torch.randn(2, 3, 4, 5)
print(f"Original shape: {tensor_4d.shape}")

# Use ellipsis for extended slicing
sliced_tensor = tensor_4d[..., 1:3, :]
print(f"Shape after slicing [..., 1:3, :]: {sliced_tensor.shape}")

# Equivalent to tensor_4d[:, :, 1:3, :]
equivalent_slice = tensor_4d[:, :, 1:3, :]
print(f"Equivalent slice shape: {equivalent_slice.shape}")

# Select first element of last dimension across all others
first_last_dim = tensor_4d[..., 0]
print(f"Shape selecting first of last dim: {first_last_dim.shape}")

Original shape: torch.Size([2, 3, 4, 5])
Shape after slicing [..., 1:3, :]: torch.Size([2, 3, 2, 5])
Equivalent slice shape: torch.Size([2, 3, 2, 5])
Shape selecting first of last dim: torch.Size([2, 3, 4])

Conclusion

Index-based operations in PyTorch provide flexible and efficient ways to access, modify, and rearrange tensor elements. From basic indexing and slicing to advanced boolean masking and ellipsis notation, these techniques enable precise data manipulation essential for deep learning workflows.