Selecting Elements from a Tensor

Selecting Elements from a Tensor in PyTorch

In PyTorch, selecting elements from a tensor works a lot like slicing arrays in NumPy—but with a few extra tricks. Whether you need a single value, a whole row/column, or specific elements based on conditions, PyTorch has you covered.

---

1. Selecting with Indexing ([])

If you’ve used NumPy before, this will feel familiar.

Let’s say we have a 2D tensor:

import torch

tensor = torch.tensor([
    [10, 20, 30],
    [40, 50, 60],
    [70, 80, 90]
])

print(tensor[2, 1])  # Single value → 80
print(tensor[:, 1])  # Second column → tensor([20, 50, 80])
print(tensor[1, :])  # Second row → tensor([40, 50, 60])

For higher dimensions, the logic is the same:

tensor[:, 2, :]  # Select all rows, 3rd column, all depth slices

---

2. Selecting with index_select

Sometimes, you want to grab specific rows or columns by index. That’s where index_select comes in.

a = torch.tensor([[1, 2], [3, 4], [5, 6]])

# Indices must be a LongTensor
indices = torch.tensor([0, 2], dtype=torch.long)

# Select rows (dim=0)
rows = torch.index_select(a, dim=0, index=indices)

# Select columns (dim=1)
cols = torch.index_select(a, dim=1, index=indices)

print(rows)  # tensor([[1, 2], [5, 6]])
print(cols)  # tensor([[1, 2], [3, 4], [5, 6]])

Key points:

• dim=0 → row selection

• dim=1 → column selection

• index must be a 1D LongTensor

---

3. Selecting with a Mask

Masks let you filter elements based on conditions.
A mask is a Boolean tensor (torch.bool) where True means “keep this element.”

a = torch.tensor([10, 20, 30, 40])
mask = torch.tensor([True, False, True, False])

selected = torch.masked_select(a, mask)
print(selected)  # tensor([10, 30])

You can also create masks dynamically:

mask = a > 25
selected = torch.masked_select(a, mask)
print(selected)  # tensor([30, 40])

Note: masked_select always returns a 1D tensor, no matter the original shape.

---

Quick Recap

• [] indexing → Grab by position (just like NumPy)

• index_select → Pick specific rows/columns with indices

• masked_select → Filter with a Boolean mask

By mastering these, you can slice, dice, and filter tensors however you want—without breaking shape logic or wasting memory.