Skip to content

← Indexing and Selection step 8 of 9

Easy Primitives

Everything above the threshold

x[x > threshold]

One expression, two things happening. x > threshold builds a boolean tensor the same shape as x; indexing with it returns a 1-D tensor of the elements where it was true, in row-major order.

The property that bites later

The output shape depends on the values, not on the shapes:

x = [1, 5, 2, 8]   threshold 3   ->  [5, 8]     length 2
x = [1, 5, 2, 8]   threshold 0   ->  [1,5,2,8]  length 4

Nothing else in this track does that. Every view and every gather has a shape you can compute from the inputs’ shapes alone, and a boolean mask does not.

Three consequences worth carrying:

  • It must allocate. There is no stride whose length depends on the contents.
  • torch.compile and CUDA graphs specialise on shapes, so a mask select forces a recompile or a graph break whenever the count changes.
  • On a GPU it needs a synchronisation, because the host has to learn the output size before it can allocate.

None of that makes it wrong. It makes it a thing to keep out of the innermost loop of a hot model, which is why masked_fill and where, whose output shape is fixed, are what attention actually uses.

Your task

def above(x: torch.Tensor, threshold: float) -> torch.Tensor

Return a 1-D tensor of the elements of x strictly greater than threshold, in row-major order, in at most three dispatched operations.

The starter loops in Python and appends to a list.