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