We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
← Indexing and Selection step 8 of 9
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.compileand 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.
Stuck?
PyTorch reference solution
Sign in to attempt this problem and reveal the reference solution.