02 · Indexing & Reshaping 🔪

4 min read

Rarely do you use a whole tensor at once. You grab rows, columns, sub-blocks — and constantly change a tensor's shape to fit the next operation.

1. The 'Why'

Real data almost never arrives in the exact shape your next operation needs. A batch of images might come in as [batch, height, width, channels] but a convolution layer wants [batch, channels, height, width]. A flat vector of 784 numbers needs to become a 28×28 image to visualize it. You'll spend a surprising amount of your PyTorch life rearranging tensors — pulling out slices, flattening, and reshaping — so fluency here pays off constantly.

The subtle part, and the reason this deserves its own sub-module, is memory. Many reshaping and slicing operations don't copy your data — they hand you a view that shares the same underlying memory as the original tensor. That's wonderfully efficient, but it means mutating the view mutates the original, which produces some of the most baffling "the value changed on its own" bugs beginners hit. Understanding when you have a view versus a copy is what separates confident PyTorch users from frustrated ones.

2. Core Concepts

Indexing and slicing work just like NumPy: t[row, col], with : selecting an entire axis and start:stop selecting a range. Negative indices count from the end.

Reshaping changes how the same elements are laid out into dimensions. reshape(rows, cols) requires the total element count to stay the same (a 12-element tensor can become 3×4 but not 3×5). Passing -1 for one dimension asks PyTorch to infer it.

Views vs. copies. .view() always returns a view sharing memory (and requires the tensor to be contiguous). .reshape() returns a view when it can and silently copies when it can't. Slicing typically returns a view too. When you need a guaranteed-independent tensor, call .clone().

3. Code in Action

Indexing and slicing

import torch

grid = torch.arange(12).reshape(3, 4)   # 3x4 matrix of values 0..11
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11]])

print(grid[0, 0])        # tensor(0)  -> top-left element [row 0, col 0]
print(grid[1])           # tensor([4, 5, 6, 7]) -> the entire second row
print(grid[:, 2])        # tensor([ 2,  6, 10]) -> the third column (all rows)
print(grid[0:2, 1:3])    # a 2x2 sub-block: rows 0-1, columns 1-2
print(grid[-1])          # tensor([ 8,  9, 10, 11]) -> the last row

Reshaping and flattening

import torch

x = torch.arange(12)     # 1-D tensor: [0, 1, ..., 11]

grid = x.reshape(3, 4)   # give the same 12 elements a 3x4 shape
flat = grid.reshape(-1)  # flatten back to 1-D; -1 lets PyTorch compute the 12

# unsqueeze/squeeze add or remove size-1 dimensions (common for batching)
row = torch.tensor([1, 2, 3])   # shape (3,)
batched = row.unsqueeze(0)      # shape (1, 3) -> add a batch dimension at front
print(batched.shape)            # torch.Size([1, 3])
print(batched.squeeze().shape)  # torch.Size([3]) -> remove size-1 dims

Views vs. copies (the gotcha)

import torch

x = torch.arange(6)      # [0, 1, 2, 3, 4, 5]

view = x.reshape(2, 3)   # usually a VIEW sharing memory with x
view[0, 0] = 99          # mutate the view...
print(x[0])              # tensor(99) -> ...and x changed too!

safe = x.reshape(2, 3).clone()  # clone() forces an independent copy
safe[0, 0] = -1                 # mutate the copy...
print(x[0])                     # tensor(99) -> x is untouched this time

4. Common Pitfalls

Pitfall 1 — Accidentally sharing memory via views. Because reshape, .view(), and slicing often return a view, mutating one tensor silently changes another. If you need independence, .clone(). Conversely, don't clone reflexively — needless copies waste memory.

Pitfall 2 — .view() on a non-contiguous tensor. After operations like .T (transpose), a tensor may be non-contiguous, and .view() will raise an error. Use .reshape() (which handles it) or call .contiguous() first.

Pitfall 3 — Reshape element-count mismatch. reshape(3, 5) on a 12-element tensor fails because 3×5 ≠ 12. When juggling shapes, print .shape liberally, and lean on -1 to let PyTorch fill in one dimension for you.

5. Further Reading & Watch List

📖 Deep Learning with PyTorch, Ch. 3 — indexing and storage
🌐 Learn PyTorch — Indexing & reshaping
🛠️ torch.reshape docs · torch.Tensor.view docs

⬅️ Prev: 01 · Creating Tensors · ➡️ Next: 03 · Tensor Math