PyTorch torch.where: masks, broadcasting, and gradients

torch.where(condition, input, other) chooses values element by element: it takes input where the Boolean condition is true and other where it is false. Use it to replace values while keeping a tensor's structure, select between two calculations, or apply a fallback to selected rows or columns.

The one-argument call, torch.where(condition), serves a different purpose: it returns index tensors. This guide tests both forms, shows how masks broadcast, and demonstrates why a finite output does not always mean a valid gradient.

All six examples below run independently on CPU with Python 3.12.14 and PyTorch 2.15.0+cpu. The accompanying checks verify values, shapes, expected failures, and gradients. They do not measure GPU performance.

Select values with a Boolean condition

Start with a comparison such as x > 0. Its result is a Boolean tensor with the same shape as x. In this example, positive entries keep their values; every other entry becomes -9.0. The second call produces integer flags, and the final expression combines two comparisons.

import torch

x = torch.tensor([[-3., 2., 0.], [4., -1., 5.]])
condition = x > 0
selected = torch.where(condition, x, -9.0)
flags = torch.where(condition, 1, 0)
between = (x > 0) & (x < 5)
print(selected.tolist())
print(flags.tolist())
print(between.tolist())

[[-9.0, 2.0, -9.0], [4.0, -9.0, 5.0]]
[[0, 1, 0], [1, 0, 1]]
[[False, True, False], [True, False, False]]

Zero follows the false branch because the comparison is strictly greater than zero. The input remains unchanged in this example. For elementwise compound conditions, put each comparison in parentheses and join them with & for AND or | for OR. Python's and and or do not perform this elementwise operation.

The torch.where API reference describes the two call forms. With three arguments, the condition must be Boolean and the selected values may come from tensors or scalars. A scalar fallback is convenient when every rejected position receives the same value.

Broadcast a mask across rows or columns

For a matrix shaped (2, 3), a mask shaped (3,) applies across columns. A mask shaped (2, 1) applies across rows. The trailing dimensions determine the alignment; the variable name does not.

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
row_mask = torch.tensor([True, False])[:, None]
fallback = torch.tensor([10., 20., 30.])
by_row = torch.where(row_mask, x, fallback)
column_mask = torch.tensor([True, False, True])
by_column = torch.where(column_mask, x, 0.0)
print(tuple(row_mask.shape), tuple(by_row.shape))
print(by_row.tolist())
print(by_column.tolist())

(2, 1) (2, 3)
[[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
[[1.0, 0.0, 3.0], [4.0, 0.0, 6.0]]

The first row is kept. The second row receives the three values in fallback. Adding [:, None] turns a length-two mask into a two-row column, making its intended alignment explicit. The column mask in the second call keeps columns zero and two in both rows.

PyTorch's broadcasting rules compare dimensions from the right. Matching sizes or a size of one can align; missing leading dimensions can be supplied. All three operands must be compatible, including other. Read the shapes together before changing them.

A square matrix can hide a mistake: a length-two mask fits a (2, 2) tensor but acts on columns. Use a non-square test fixture when checking row masks. In a data preparation pipeline, confirm which axis represents samples before passing tensors into a PyTorch DataLoader.

Return indices with torch.where(condition)

For this two-dimensional input, the one-argument form returns two tensors: row indices and column indices. Read corresponding entries as coordinate pairs. Passing that tuple back into the input extracts the matching values.

import torch

x = torch.tensor([[2, -1, 4], [0, 5, -2]])
condition = x > 0
rows, columns = torch.where(condition)
values = x[rows, columns]
filled = torch.where(condition, x, 0)
empty_indices = torch.where(x > 10)
print(rows.tolist(), columns.tolist())
print(values.tolist(), tuple(values.shape))
print(filled.tolist(), tuple(filled.shape))
print([part.tolist() for part in empty_indices])

[0, 0, 1] [0, 2, 1]
[2, 4, 5] (3,)
[[2, 0, 4], [0, 5, 0]] (2, 3)
[[], []]

The coordinates are (0, 0), (0, 2), and (1, 1). The extracted values form a one-dimensional tensor of length three. The three-argument call keeps a two-by-three matrix and puts zeros in rejected positions.

If nothing matches, this matrix still produces two index tensors, both empty. Choose the return form based on what the next step needs: coordinates, a compact selection of values, or a tensor with replacement values. The one-argument form is equivalent to torch.nonzero(condition, as_tuple=True); it does not return the selected data itself.

Diagnose dtype and shape errors

These negative examples deliberately fail. They print exception classes rather than full messages, whose wording can vary between releases.

import torch

x = torch.ones(2, 3)
cases = {
    "integer condition": lambda: torch.where(
        torch.tensor([1, 0, 1]), x, 0.0),
    "row mask without axis": lambda: torch.where(
        torch.tensor([True, False]), x, 0.0),
    "incompatible other": lambda: torch.where(
        x > 0, x, torch.zeros(2, 4)),
}
errors = {}
for name, operation in cases.items():
    try:
        operation()
    except RuntimeError as exc:
        errors[name] = type(exc).__name__
        print(name + ": " + errors[name])
    else:
        raise AssertionError("Expected failure: " + name)

integer condition: RuntimeError
row mask without axis: RuntimeError
incompatible other: RuntimeError

The first case supplies integer flags where this call requires a Boolean mask. Build the condition from a comparison that expresses your rule, such as flags != 0. The second case tries to align a length-two vector with a trailing dimension of three; add the intended row axis as shown earlier. The last case gives the fallback four columns while the input has three.

Do not reshape tensors just until the exception disappears. Write down what each dimension means, decide whether the condition applies to samples or features, and check the resulting values on a small matrix. A compatible shape can still implement the wrong selection.

Gradients follow selected finite values

Here the Boolean mask is fixed. Each selected input value contributes its weight to the loss. The scalar fallback is reused at three positions, so its gradient adds those positions' weights.

import torch

condition = torch.tensor([[True, False, True],
                          [False, False, True]])
x = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                 dtype=torch.float64, requires_grad=True)
fallback = torch.tensor(10., dtype=torch.float64, requires_grad=True)
weights = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                       dtype=torch.float64)
selected = torch.where(condition, x, fallback)
loss = (selected * weights).sum()
loss.backward()
print(selected.tolist())
print(loss.item())
print(x.grad.tolist())
print(fallback.grad.item())

[[1.0, 10.0, 3.0], [10.0, 10.0, 6.0]]
156.0
[[1.0, 0.0, 3.0], [0.0, 0.0, 6.0]]
11.0

The input gradient is zero at the rejected positions and equals the weight at selected positions. The fallback gradient is 2 + 4 + 5 = 11. The test suite checks these results against an independent scalar calculation and central finite differences.

This example does not learn the Boolean decision. A comparison supplies a discrete mask; do not expect a trainable threshold to receive a useful gradient through that comparison. If you only need inference and want to control gradient recording, see the separate guide to eval, no_grad, and inference_mode.

Prevent invalid arithmetic before selecting the result

Ordinary Python evaluates the arguments before calling torch.where. In torch.where(valid, x / denominator, 0.0), the division has already happened. Selecting zero afterward can conceal an invalid operation in the output while leaving a NaN gradient.

import torch

denominator = torch.tensor([0., 2.])
valid = denominator != 0
unsafe_x = torch.tensor([3., 6.], requires_grad=True)
unsafe = torch.where(valid, unsafe_x / denominator, 0.0)
unsafe.sum().backward()
print(unsafe.tolist(), torch.isfinite(unsafe_x.grad).tolist())

safe_x = torch.tensor([3., 6.], requires_grad=True)
safe_denominator = torch.where(valid, denominator, 1.0)
safe = torch.where(valid, safe_x / safe_denominator, 0.0)
safe.sum().backward()
print(safe.tolist(), safe_x.grad.tolist())

[0.0, 3.0] [False, True]
[0.0, 3.0] [0.0, 0.5]

Both forward results are [0.0, 3.0]. The first printed finiteness check shows that the unsafe gradient has a non-finite entry. The second calculation replaces zero denominators with one before division, then explicitly selects zero at those positions. Its gradient is [0.0, 0.5].

This fixes the demonstrated zero-denominator case with finite inputs. It does not validate arbitrary NaNs, infinities, logarithms, or square roots. Each operation needs valid operands for the path recorded by autograd. PyTorch's autograd notes explain why masking an invalid result after division can still damage gradients.

Which selection operation should you use?

• Use torch.where(mask, x, fallback) to choose between values across a broadcasted output.
• Use torch.where(mask) when you need coordinates.
• Use Boolean indexing such as x[mask] to extract matching elements from the same-shaped matrix shown here.
• Use an integer-index operation such as torch.gather when positions are already specified numerically along an axis.

Before using a selection in a training loop, check its output shape, a few exact values, and the gradients of any branch that performs division or another restricted-domain operation. Those checks catch errors that a plausible-looking output can miss.