Indexing and Slicing PyTorch Tensors

PyTorch tensors support the same indexing syntax as NumPy: integers pick a position, slices pick a range, boolean masks pick elements that satisfy a condition, and lists or tensors of integers pick arbitrary positions. Mastering these lets you pull out batches, rows, columns and filtered values without writing loops.

Basic indexing and slicing

import torch

x = torch.arange(12).reshape(3, 4)
print(x)
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])

print(x[0])         # first row: tensor([0, 1, 2, 3])
print(x[1, 2])      # tensor(6)
print(x[:, 1])      # second column: tensor([1, 5, 9])
print(x[1:, :2])    # rows 1.., columns 0-1
print(x[::2])       # every second row
print(x[-1])        # last row

The syntax is start:stop:step per dimension, and the stop value is excluded. A single integer removes that dimension, so x[0] is 1-D. To keep the dimension, slice instead: x[0:1] has shape (1, 4).

Slices are views

Basic slices share memory with the original tensor. Changing the slice changes the source.

row = x[0]
row[0] = 100
print(x[0, 0])      # tensor(100)

safe = x[0].clone() # independent copy

Use .clone() when you need a separate copy.

Boolean masks

A comparison produces a bool tensor, which you can use as an index. The result is always a flattened 1-D tensor of the selected elements.

t = torch.tensor([3, -1, 4, -5, 9])
mask = t > 0
print(mask)             # tensor([ True, False,  True, False,  True])
print(t[mask])          # tensor([3, 4, 9])

t[t < 0] = 0            # assign through a mask
print(t)                # tensor([3, 0, 4, 0, 9])

This is how you clip values, ignore padding tokens or select samples of one class. Combine conditions with & and |, and wrap each condition in parentheses: (t > 0) & (t < 5).

Integer index lists

x = torch.arange(12).reshape(3, 4)
rows = torch.tensor([0, 2])
print(x[rows])              # rows 0 and 2

r = torch.tensor([0, 1, 2])
cidx = torch.tensor([3, 0, 1])
print(x[r, cidx])           # tensor([3, 4, 9]): pairs (0,3), (1,0), (2,1)

When you pass two index tensors, PyTorch pairs them up element by element. This "advanced indexing" returns a copy, not a view.

torch.where, index_select and gather

a = torch.tensor([1.0, -2.0, 3.0])
print(torch.where(a > 0, a, torch.zeros_like(a)))   # tensor([1., 0., 3.])

x = torch.arange(12).reshape(3, 4)
print(torch.index_select(x, dim=1, index=torch.tensor([0, 3])))

# gather: pick one value per row using per-row column indices
idx = torch.tensor([[3], [0], [1]])
print(x.gather(1, idx))     # tensor([[3], [4], [9]])

torch.where(cond, a, b) chooses element-wise between two tensors. index_select picks whole slices along one dimension. gather is useful in classification: given predicted scores of shape (batch, classes) and a column of labels, it fetches the score of the correct class for each sample.

Ellipsis and None

img = torch.rand(8, 3, 32, 32)
print(img[..., 0].shape)        # last dim index 0 -> (8, 3, 32)
v = torch.tensor([1, 2, 3])
print(v[:, None].shape)         # (3, 1): adds a dimension
print(v[None, :].shape)         # (1, 3)

... stands for "all remaining dimensions", and None inserts a new dimension of size one, which is very useful with broadcasting (next post).

Putting indexing to work

Indexing becomes much more useful once you apply it to realistic shapes. Suppose a batch of sentences is stored as a tensor of token ids with shape (batch, seq_len). The expression tokens[:, 0] returns the first token of every sentence, and tokens[:, -1] returns the last one. If padding uses id 0, then tokens != 0 gives a mask that marks the real tokens, and summing it along dimension 1 gives the true length of each sentence.

The same ideas apply to images. Given a batch of shape (batch, channels, height, width), the slice imgs[:, 0] selects the first channel of every image, and imgs[:, :, :16, :16] crops the top-left corner of all images at once. Because basic slices are views, cropping is almost free, but remember that writing into the crop modifies the original batch.

Boolean masks pair naturally with assignment. In a loss computation you might set ignored positions to zero with loss[mask == 0] = 0, or replace invalid scores with a large negative number before a softmax. This is the exact technique used for attention masks in transformers, so it is worth practising now on small tensors where you can print every step and confirm the result by eye.

Common mistakes

• Expecting a mask result to keep its shape: it is flattened.
• Forgetting that slices are views and accidentally modifying the original data.
• Writing t > 0 & t < 5 without parentheses, which fails because of operator precedence.
• Using float tensors as indices: index tensors must be integer (long) or bool.

Practise by selecting the last column of a matrix, all rows whose first value is positive, and the value at the argmax of each row.