Indexing and Slicing

Indexing is everything you remember from NumPy plus a few extras

If you wrote NumPy comfortably, the next ten minutes are revision. The PyTorch indexing rules are:

Basic indexing — integer and slice notation. Returns a view when possible (no copy).
Boolean indexing — index with a bool tensor of matching shape. Returns a 1-D tensor of selected elements (always a copy).
Advanced (fancy) indexing — index with an integer tensor. Always a copy.
Mixed — combinations of the above; PyTorch matches NumPy's rules.

The single most important fact: basic slicing returns a view, not a copy. If you write into the slice, you write into the original. That's a feature, not a bug — it makes in-place updates cheap — but it's also the source of the "I changed one tensor and three others moved" debugging story.

The selectors you'll use the most

x[:, 0] — first column of every row.
x[..., -1] — last element on the last dim, regardless of how many dims x has. Beautifully concise.
x[None, :] / x.unsqueeze(0) — add a batch dim.
x[mask] — boolean filter.
torch.gather(x, dim, index) — pick values along a dim using an index tensor. Hugely common in attention and loss code.

Code

Basic indexing and slicing·python

import torch

t = torch.tensor([
    [ 1,  2,  3,  4],
    [ 5,  6,  7,  8],
    [ 9, 10, 11, 12],
])

t[0]          # tensor([1, 2, 3, 4])  — first row
t[0, 2]       # tensor(3)             — single element
t[-1]         # last row
t[:, 1]       # all rows, col 1 → tensor([2, 6, 10])
t[0:2, :]     # first two rows
t[::2, :]     # every other row
t[..., -1]    # last col, regardless of rank → tensor([4, 8, 12])

Boolean and fancy indexing·python

import torch

t = torch.tensor([10, 20, 30, 40, 50])

# Boolean — same-shape bool tensor
mask = t > 25
t[mask]                  # tensor([30, 40, 50])

# Fancy — integer index tensor
idx = torch.tensor([0, 2, 4])
t[idx]                   # tensor([10, 30, 50])

# Mix in 2D
m = torch.arange(20).reshape(4, 5)
m[m % 3 == 0]            # all multiples of 3, flattened
m[[0, 2, 3]]             # rows 0, 2, 3
m[[0, 2, 3], [1, 1, 1]]  # element-wise: (0,1), (2,1), (3,1)

View vs copy — the gotcha·python

import torch

a = torch.arange(12).reshape(3, 4)
row = a[1]               # VIEW into a (no copy)
row[:] = 99              # writes into a too!
print(a)
# tensor([[ 0,  1,  2,  3],
#         [99, 99, 99, 99],
#         [ 8,  9, 10, 11]])

# To get an independent copy:
row_copy = a[1].clone()
row_copy[:] = -1
print(a[1])              # untouched: tensor([99, 99, 99, 99])

torch.gather — the attention/loss workhorse·python

import torch

# Pick the predicted-class probability for each sample
logits = torch.randn(4, 5).softmax(-1)   # (batch=4, classes=5)
labels = torch.tensor([2, 0, 4, 1])

# We want logits[i, labels[i]] for each i
chosen = logits.gather(1, labels.unsqueeze(1)).squeeze(1)
print(chosen.shape)  # torch.Size([4])

# Equivalent (less efficient) loop:
# chosen = torch.tensor([logits[i, labels[i]] for i in range(4)])

Exercise

Build a (4, 5) tensor of logits with torch.randn. Given a label tensor [2, 0, 4, 1], extract the logit for the labeled class of each sample using BOTH a Python loop AND torch.gather. Verify they produce the same values. Use timeit to compare on a (1024, 1000) tensor — gather should be dramatically faster.

Indexing is everything you remember from NumPy plus a few extras

The selectors you'll use the most

Code

External links

Exercise

Progress

Comments 0