Indexing and slicing tensors
Why this matters
Every PyTorch model needs to access subsets of data: selecting a batch, extracting specific features, filtering predictions. Inefficient indexing becomes a bottleneck in data loading and model inference.
Explanation
Tensor indexing retrieves elements or subarrays by position, range, or condition. Like Python lists, tensors support integer indexing, slicing, and boolean masking: but with important GPU memory implications. How it works mechanically: Integer indices select single elements or preserve dimensions, slice notation (start:stop:step) creates views of contiguous data, and boolean masks filter rows based on conditions. The key difference from NumPy: indexing may move data between GPU and CPU if you mix tensor types, and advanced indexing (fancy indexing) sometimes creates copies instead of views. When to use: Integer and slice indexing for selecting batches, features, or time steps; boolean indexing for filtering based on conditions or masks (e.g., selecting high-confidence predictions).
Analogy
Indexing a tensor is like labeling a shelf: [0] grabs the first item, [1:3] grabs items 1 and 2, [tensor > 5] grabs all items matching your label condition.
Code
import torch
tensor = torch.arange(12).reshape(3, 4)
print("Original tensor:")
print(tensor)
print()
print("Single element (row 1, col 2):")
print(tensor[1, 2])
print()
print("Entire row (row 0):")
print(tensor[0])
print()
print("Slice rows 0-1, all columns:")
print(tensor[0:2])
print()
print("Slice rows, specific columns (all rows, cols 1-3):")
print(tensor[:, 1:3])
print()
print("Boolean mask (elements > 6):")
mask = tensor > 6
print(mask)
print()
print("Extract values matching mask:")
filtered = tensor[mask]
print(filtered)
print()
print("Negative indexing (last row):")
print(tensor[-1])
print()
print("Step in slicing (every other column):")
print(tensor[:, ::2])Original tensor:
tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
Single element (row 1, col 2):
tensor(6)
Entire row (row 0):
tensor([0, 1, 2, 3])
Slice rows 0-1, all columns:
tensor([[0, 1, 2, 3],
[4, 5, 6, 7]])
Slice rows, specific columns (all rows, cols 1-3):
tensor([[ 1, 2, 3],
[ 5, 6, 7],
[ 9, 10, 11]])
Boolean mask (elements > 6):
tensor([[False, False, False, False],
[False, False, False, True],
[ True, True, True, True]])
Extract values matching mask:
tensor([ 7, 8, 9, 10, 11])
Negative indexing (last row):
tensor([ 8, 9, 10, 11])
Step in slicing (every other column):
tensor([[ 0, 2],
[ 4, 6],
[ 8, 10]]) What just happened?
Created a 3×4 tensor, then extracted elements six different ways: (1) single value by row-column index, (2) entire row by row index alone, (3) row range using slice notation, (4) column range using : and slice notation, (5) boolean mask filtering values > 6 (returns 1D), (6) negative indexing for the last row, (7) step slicing for every other column. Each operation returned either a scalar, 1D, or 2D tensor depending on how many dimensions were indexed.
Common gotcha
Boolean indexing flattens the result to 1D, even if your original tensor is 2D or higher. If you need to keep dimensions, use torch.where() or explicit row/column indexing instead. Also, slicing creates a view (shares memory) but boolean indexing creates a copy: so modifying filtered values won't affect the original.
Error recovery
IndexError: index X is out of bounds for dimension Y with size ZRuntimeError: tensors used as indices must be long, byte or bool tensorsTypeError: only integers, slices, ellipsis, torch.newaxis and torch.Tensor objects are valid indicesExperienced dev note
Boolean indexing is tempting but creates a copy and flattens dimensions: this kills performance in data loading loops. For filtering batches, use gather() or advanced indexing with torch.where() to keep your shape. Also: never index inside torch.compile()'s traced function with Python control flow (if tensor[i] > 5:), because the trace captures one specific path. Use tensor operations instead.
Check your understanding
If you create a 2D tensor of shape (100, 50) and apply boolean indexing tensor[tensor > 0.5], what is the shape of the result, and why did the shape change compared to slicing tensor[0:50, 10:20]?
Show answer hint
Boolean indexing flattens to 1D because the condition is applied element-wise and only matching elements are returned. Slicing preserves dimensions because it selects contiguous ranges along each axis.