Indexing and Slicing

Indexing selects elements by position. The rule that trips people up: a slice keeps the dimension it indexes, while a single integer drops it. Internalize that and most shapes become predictable.

The idea

• t[i] selects along dimension 0. On a 2-D tensor, t[0] is a row of shape (n,).
• t[i, j] selects one scalar as a 0-dim tensor; .item() converts it to a Python number.
• t[start:stop:step] slices. Omitted bounds mean "from the beginning" or "through the end".
• t[:, 0] selects column 0 across all rows → (m,), since the integer index drops that axis.
• t[:2, 1:] is a sub-block; the colon keeps each axis.
• t[t > 0] is boolean masking, returning a 1-D tensor of the selected values and discarding layout.
• t[..., -1] uses ... for "all remaining leading axes".

Slices are views in PyTorch, so they share storage with the base tensor. Writing into a slice mutates the original; .clone() detaches a copy. For data-dependent lookups, torch.gather and torch.index_select collect values along an axis and stay differentiable.

Worked example

Let x = torch.arange(12).reshape(3, 4):

• x[1] → (4,), [4, 5, 6, 7]
• x[:, 1] → (3,), [1, 5, 9]
• x[1:3, 2:] → (2, 2), [[6, 7], [10, 11]]
• x[-1, -1] → 11
• x[x % 3 == 0] → tensor([0, 3, 6, 9])

For a batch of images (32, 3, 224, 224): images[0] is one image (3, 224, 224), images[:8] is a mini-batch (8, 3, 224, 224), and images[:, 0] is the red channel of every image (32, 224, 224) because the byte index dropped the channel axis.

In code

import torch

x = torch.arange(12).reshape(3, 4)
print(x[1].shape)          # torch.Size([4])
print(x[:, 1].shape)       # torch.Size([3])
print(x[1:3, 2:])          # tensor([[ 6,  7], [10, 11]])
print(x[-1, -1].item())    # 11
print(x[x % 3 == 0])       # tensor([0, 3, 6, 9])

row = x[0]                 # a view, not a copy
row[0] = -1
print(x[0, 0].item())      # -1

Check yourself

1. For an image batch (32, 3, 224, 224), what is the shape of images[:, 1] and why?
2. Why does x[1] lose the first axis while x[1:2] keeps it?
3. What does x[x > 0] = 0 do, and what shape does x have afterward?

Key takeaways

• Integer indices drop an axis; slices keep it.
• Boolean masks return a flat 1-D tensor of the selected values.
• Slices are views — clone before writing if you need to preserve the original.