Table of Contents
Indexing Tensors
Indexing is how you select elements from a tensor. In PyTorch, indexing looks similar to NumPy and Python lists, but it works across any number of dimensions. If x is a tensor, x[i] selects along the first dimension. For a 2D tensor, x[i, j] selects a single element at row i and column j. For higher dimensions, you keep adding indices, for example x[n, c, h, w] for a typical image batch layout.
When you index a tensor with fewer indices than it has dimensions, PyTorch keeps the remaining dimensions. For example, if x has shape (10, 3, 32, 32), then x[0] has shape (3, 32, 32). This behavior is essential when working with batches because selecting one batch item removes the batch dimension.
PyTorch uses zero based indexing, so the first element is at index 0, and negative indices count from the end, so x[-1] is the last element along that dimension.
Slicing and Keeping Dimensions
Slicing selects ranges using the start:stop:step syntax. For a 1D tensor, x[2:5] returns elements with indices 2, 3, 4. The stop index is exclusive. Omitting start means start at the beginning, omitting stop means go to the end, and omitting step means step size 1. For example, x[:5], x[5:], and x[::2] are all valid.
For multi dimensional tensors, you can slice each dimension separately, for example x[:2, :, 10:20]. The colon : means select everything along that dimension. This is a common pattern for slicing batches, channels, and spatial regions.
A key detail is whether indexing removes a dimension. Using an integer index removes that dimension, while using a slice keeps it. For a 2D tensor x with shape (N, D), x[0] has shape (D,), but x[0:1] has shape (1, D). Keeping dimensions is often helpful when you want shapes to remain compatible with later operations.
Important rule: an integer index drops that dimension, a slice keeps that dimension.
If you specifically want to keep a dimension while selecting a single index, you can use None (or unsqueeze which is covered elsewhere). For example, x[0, None, :] inserts a size 1 dimension in the middle.
Advanced Indexing: Lists, Tensors, and Masks
Beyond slices, PyTorch supports selecting with index lists or index tensors. If idx = torch.tensor([0, 3, 7]), then x[idx] gathers those positions along the first dimension. This is common when you want a subset of batch items.
Boolean masking selects all elements where a condition is true. If mask = x > 0, then x[mask] returns a 1D tensor of all positive elements, flattened because it is selecting individual entries, not preserving the original grid. Masking is very useful for filtering, computing statistics on a subset, or ignoring padded values.
Be careful that advanced indexing can change shapes in ways that are less intuitive than slicing. Slicing tends to preserve structure, while boolean masks and index arrays tend to produce gathered results whose shape follows the index objects.
Important rule: boolean masking like x[mask] typically flattens the selected elements into a 1D result, it does not preserve the original shape.
Views, Copies, and In Place Changes
Many slices in PyTorch return a view into the same underlying storage, meaning they share memory. If you modify the slice, you can modify the original tensor. This is fast and memory efficient, but it can surprise you during debugging. Advanced indexing, like using a list of indices or a boolean mask, usually returns a copy instead of a view.
If you need to be sure you have an independent tensor, you can call clone() on the result.
Important rule: slices often share memory with the original tensor, advanced indexing often creates a copy. If you plan to modify the result safely, consider using clone().
Broadcasting: The Shape Matching Superpower
Broadcasting lets PyTorch apply elementwise operations to tensors with different shapes by implicitly expanding size 1 dimensions. This is why you can add a vector to every row of a matrix, or subtract per channel means from an image tensor, without writing loops.
The core idea is that when doing an elementwise operation like addition or multiplication, PyTorch compares shapes from the rightmost dimension to the left. Two dimensions are compatible if they are equal, or one of them is 1. If compatible, the dimension of size 1 is virtually expanded to match the other size.
For example, if A has shape (N, D) and b has shape (D,), then A + b works because (D,) is treated like (1, D), and the leading 1 can expand to N, producing a result of shape (N, D).
More generally, given shapes (a_1, a_2, ..., a_k) and (b_1, b_2, ..., b_m), you align them to the right by pretending the shorter shape has leading ones, then check each dimension pair for equality or 1. The output dimension is the maximum of the two.
Broadcasting rule: for each dimension from the right, sizes must match or one must be 1. Otherwise you get a runtime error.
Broadcasting in Practice with Common Deep Learning Shapes
Broadcasting is especially common with batch dimensions. If x has shape (B, D) and scale has shape (1, D), then x * scale scales features per column for every batch item. If bias has shape (D,), it also broadcasts across the batch.
For images, if x has shape (B, C, H, W) and you have per channel means mean with shape (C,), you typically reshape it to (1, C, 1, 1) so it broadcasts across batch, height, and width. The important part is placing 1 in the dimensions you want to expand.
Broadcasting is also useful for pairwise computations. If you have a with shape (N, 1) and b with shape (1, M), then a - b produces an (N, M) matrix of differences.
Debugging Shape Errors from Indexing and Broadcasting
Most beginner errors here are shape mismatches that only show up at runtime. A reliable approach is to print or inspect shapes right before an operation that fails. If broadcasting fails, check shapes from the rightmost dimension and locate the first incompatible pair.
A second common issue is accidentally dropping a dimension by using an integer index, then later expecting it to still exist. If you intended to keep it, switch from x[i] to x[i:i+1] or insert a singleton dimension.
Finally, remember that boolean masking changes the structure of your tensor. If you mask and then try to reshape or combine with an unmasked tensor, you may need a different approach than x[mask], such as using masked_fill or multiplying by a mask cast to the right dtype, depending on the goal.
If a later operation expects a batch dimension, do not accidentally remove it with integer indexing like x[0]. Use slicing like x[0:1] to preserve the batch dimension.