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

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.