Everything above the threshold

A boolean mask selects the elements whose condition is true. When the mask has the same shape as the input, the result is a one-dimensional tensor.

import torch

x = torch.tensor([[1.0, 5.0], [2.0, 8.0]])
print(x[x > 3.0].tolist())  # [5.0, 8.0]
print(x[x > 0.0].tolist())  # [1.0, 5.0, 2.0, 8.0]
print(x[x > 10.0].tolist()) # []

The values appear in logical row-major order, as if you visited each row from left to right. A value equal to the threshold is excluded because the comparison is strictly greater than.

The output length depends on the values

The input shape is unchanged between these examples, but the number of selected elements differs. Reading with advanced boolean indexing creates a new tensor rather than a view sharing the input's storage.

A variable output length can complicate compiled execution and CUDA graph capture. Behavior depends on the compiler, configuration and backend; a changing count does not universally force a recompile. On CUDA, determining the output size can also require synchronization. Measure the actual workload before drawing performance conclusions from the expression alone.

If you want to keep the original shape and replace unwanted entries, consider where or masked_fill instead. Those perform a different task: they replace values rather than returning a shorter selection.

Select values using coordinates

A slice selects a regularly spaced region of an array. A gather selects entries using supplied indices. Check whether those indices select whole rows or columns, or a different entry within each row.

For [[10, 11], [20, 21]], choosing column 1 from the first row and column 0 from the second gives [11, 20]. Choosing columns [1, 0] from every row instead gives [[11, 10], [21, 20]]. These operations use the same index values but have different meanings and output shapes.

Scatter writes values to selected destinations

Scatter takes values and destination indices. The operation may replace existing values or combine them using a reduction such as addition.

Starting with [0, 0, 0], scatter-add of values [2, 3] at indices [1, 1] gives [0, 5, 0]. Both values contribute to the repeated destination. Repeated assignment is not equivalent to this sum; follow the operation's rule for repeated indices.

Selecting with a mask differs from replacing with a mask

For values [10, 20, 30] and mask [True, False, True], selecting masked entries gives [10, 30]. Replacing masked entries with zero gives [0, 20, 0]. Selection can change the shape; replacement preserves it. A selection whose size depends on runtime values needs particular care in compiled code with fixed output shapes.

Check whether an update changes the original array

JAX indexed updates return a new array value. In PyTorch, basic slicing can return a view that shares storage, while advanced indexing returns a copy. Modifying a selected copy does not modify the source. Direct indexed assignment is a separate operation, so distinguish changing a selection from assigning into the original.