Learn AI Create the Future — with Pranay
Main site

PyTorch Tensors

Indexing and boolean masks

Square brackets select rows, columns, windows and — with a boolean mask — every element passing a test, all without a single Python loop.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Indexing means picking out part of a tensor — one row, one column, or every element that passes a test.

Think of cleaning rice on a steel plate. You do not inspect each grain with tweezers, one at a time. Your eye marks all the dark stones at once, and your hand sweeps them away together. A mask works like that marking: it flags many elements in one look, so one sweep can act on all of them.

Why it exists

Real data always needs selecting. Keep the readings above a limit. Cap the crazy values. Grab the second column. Take the first hundred rows for testing.

You could loop element by element, asking "this one? this one?". In Python that is painfully slow, and the code buries the intent. Indexing states the what and lets PyTorch handle the each.

How it works

A mask is a tensor of yes/no answers, the same shape as your data. Comparing makes one; indexing with it keeps the yes positions.

temps:        [31, 34, 29, 40, 38, 25]
temps > 35 :  [no, no, no, yes, yes, no]     <- the mask
temps[mask]:  [40, 38]                        <- kept

For grids, positions in square brackets are read one per dimension: grid[1, 2] means row 1, column 2. A colon means "everything in this direction": grid[:, 2] is all rows, column 2.

A real example you have seen

Every spreadsheet filter is a mask. "Show orders above ₹500" marks the passing rows and hides the rest. Excel does with menus what PyTorch does with brackets — the underlying idea is identical.

Remember this

  • A mask is a same-shaped tensor of yes/no, made by a comparison.
  • data[mask] keeps the yes elements; data[mask] = v overwrites them.
  • In grids, brackets read one slot per dimension, and : means "all of it".

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Outputs verified with torch 2.5.1, CPU.

Masks, slices and windows

indexing.py
import torch

temps = torch.tensor([31., 34., 29., 40., 38., 25.])
hot = temps > 35              # a boolean mask: True where the test passes
print(hot)
print(temps[hot])             # keep only the True positions
temps[temps > 35] = 35.       # cap the readings in one line
print(temps)

grid = torch.arange(12).reshape(3, 4)
print(grid[1])                # one row
print(grid[:, 2])             # one column
print(grid[0:2, 1:3])         # a 2x2 window

# a slice is a view; a mask result is a copy
window = grid[0:2]
window[0, 0] = 100            # changes grid
picked = grid[grid > 5]
picked[0] = -1                # does NOT change grid
print(grid[0, 0].item(), grid[1, 3].item())
Output
tensor([False, False, False,  True,  True, False])
tensor([40., 38.])
tensor([31., 34., 29., 35., 35., 25.])
tensor([4, 5, 6, 7])
tensor([ 2,  6, 10])
tensor([[1, 2],
        [5, 6]])
100 7

The walkthrough

temps[temps > 35] = 35. is the idiom to internalise. Compare, select, assign — one line, no loop. This exact pattern cleans sensor glitches, clamps outliers and zeroes padding in real codebases daily.

temps[hot] came back with 2 elements. A mask selection's length depends on the data, not the shape. You cannot know it before looking — a fact with real consequences on GPUs (researcher block).

The view/copy split at the end is the trap. Plain slices like grid[0:2] are views — windows onto the same memory, so writing through window changed grid. Mask (and index-list) selections are copies — writing into picked left grid alone. Editing a selection and expecting the original to change, or not to, causes bugs in both directions. The demo prints 100 7: the slice-write landed, the mask-write did not.

Counting with a mask is a one-liner worth knowing: (temps > 35).sum() — True counts as 1. Related: mask.any(), mask.all(), and torch.nonzero(mask) for the positions themselves.

Combining conditions

Python's and/or do not work element-wise. Use &, |, ~, and bracket every comparison — & binds tighter than >:

python
temps = torch.tensor([31., 34., 29., 40., 38., 25.])
mild = temps[(temps > 30) & (temps < 36)]
print(mild)
Output
tensor([31., 34.])

Common mistakes

and instead of &. (a > 0) and (b > 0) raises Boolean value of Tensor with more than one element is ambiguous. The keywords ask for one single yes/no; masks hold many. Symbols, with brackets, always.

Forgetting the brackets. temps > 30 & temps < 36 parses as temps > (30 & temps) < 36 — a precedence accident that errors at best, lies at worst.

Assigning into a copy. grid[grid > 5][0] = -1 modifies a temporary copy that vanishes: the original never changes and nothing warns you. Direct mask assignment (grid[grid > 5] = -1) works because it is a single fused operation, not a chain.

Mask shape mismatch. A mask must match the shape (or a broadcastable prefix) of what it indexes. Masking a (3, 4) grid with a (4,) row-mask selects columns in a way few people predict — print shapes before trusting it.

Try it yourself

Make a (5, 3) tensor of random marks with torch.randint(0, 100, (5, 3)). In one line each: count how many marks are below 40; replace them with 40; then select every row whose total exceeds 200 — that last one needs a mask built from a reduction, indexed on dim 0.

What to learn next

Researcher — Mathematics and papers.

Basic versus advanced indexing

PyTorch inherits NumPy's two-regime semantics:

  • Basic indexing — integers, slices, ellipsis, None: expressible as pure stride arithmetic (offset, shape, stride adjustments), hence always a view, O(1).
  • Advanced indexing — integer tensors/lists or boolean masks: general gather, not expressible as strides, hence always a copy, O(size of selection).

Assignment forms (x[idx] = v) are handled specially in both regimes as in-place scatter into the original storage — which is why grid[grid > 5] = -1 mutates while grid[grid > 5][0] = -1 does not: the second desugars to index(...) (copy) followed by a write into the temporary.

Boolean mask indexing is internally x[mask.nonzero(...)] — mask → index conversion, then gather.

The data-dependent shape problem

x[mask] has shape determined by mask.sum(), known only at runtime. Two costs follow:

  1. Device synchronisation. On GPU, allocating the output requires the count on the host, forcing a device-to-CPU sync that stalls the pipeline. In hot loops, prefer arithmetic masking — x * mask, torch.where(mask, x, fallback) — which keeps shapes static and stays asynchronous.
  2. Compilation. Data-dependent shapes are the canonical graph-break for torch.compile and export paths; torch.where and masked arithmetic trace cleanly, x[mask] does not.

masked_fill, masked_select, index_put_ are the explicit named forms; torch.where(cond, a, b) is the differentiable branch — gradients flow through the selected side, zero through the other.

Gradient behaviour

Indexing is linear, so autograd handles it exactly: the backward of a gather is a scatter-add of the output gradient into the selected positions, zeros elsewhere. Repeated indices therefore accumulate gradient — consistent with the multivariate chain rule, and occasionally surprising when an index appears twice.

References

  • Harris et al. (2020), Array programming with NumPy, Nature 585 — the indexing semantics PyTorch adopted, §"Indexing".
  • PyTorch docs, Tensor Indexing API — the authoritative basic/advanced classification.
  • He et al. (2015), Delving Deep into Rectifiers — masking as architecture: ReLU itself is x * (x > 0), the same mask machinery inside a nonlinearity.

What to learn next