Python / PyTorch Fundamentals Interview Questions
How does PyTorch handle multi-dimensional indexing and slicing of tensors?
PyTorch tensor indexing follows NumPy-style conventions, including basic slicing, advanced (fancy) indexing with integer/boolean tensors, and the powerful ... (ellipsis) operator for indexing high-dimensional tensors concisely.
import torch x = torch.arange(24).reshape(2, 3, 4) # shape (2, 3, 4) # ââ Basic slicing â same as Python lists/NumPy print(x[0]) # shape (3, 4) â first "batch" print(x[0, 1]) # shape (4,) â first batch, second row print(x[0, 1, 2]) # scalar â single element print(x[:, 0, :]) # shape (2, 4) â all batches, first row, all cols print(x[..., 0]) # shape (2, 3) â ellipsis: all leading dims, last dim index 0 print(x[0:1, :, -1]) # shape (1, 3) â slice + negative index # ââ Boolean (mask) indexing mask = x > 10 print(x[mask]) # 1D tensor of all elements > 10 x_clamped = x.clone() x_clamped[x_clamped > 10] = 0 # zero out values > 10 # ââ Fancy (advanced) integer indexing idx = torch.tensor([0, 2]) print(x[:, idx, :]) # shape (2, 2, 4) â select specific indices along dim 1 # ââ torch.gather: select elements using an index tensor scores = torch.tensor([[0.1, 0.7, 0.2], [0.3, 0.3, 0.4]]) # (2, 3) top_idx = scores.argmax(dim=1, keepdim=True) # (2, 1) top_val = scores.gather(dim=1, index=top_idx) # (2, 1) print(top_val) # tensor([[0.7], [0.4]]) # ââ torch.where: conditional element selection result = torch.where(x > 10, x, torch.zeros_like(x)) # keep if >10, else 0 # ââ Important: most slicing returns a VIEW, not a copy! y = x[0] y[0, 0] = 999 print(x[0, 0, 0]) # 999 â x was modified too! (shared memory) # Use x[0].clone() to get an independent copy
| Pattern | Example | Returns |
|---|---|---|
| Basic slicing | x[:, 0] | View (shares memory) |
| Boolean mask | x[x > 0] | Copy (1D, new memory) |
| Fancy indexing | x[:, [0,2]] | Copy (new memory) |
| Ellipsis | x[..., 0] | View — skips middle dims |
| gather | x.gather(dim, index) | Copy — selects per index |
More Related questions...