Skip to content

← Indexing and Selection step 7 of 9

Easy Primitives

Rows 2 through 5, cheaply

Two ways to take some rows, and the difference is whether the rows you want happen to be next to each other.

x.index_select(0, torch.tensor([2, 3, 4, 5]))   # a copy, always
x.narrow(0, start=2, length=4)                  # a view, always
x[2:6]                                          # the same view

index_select accepts any index at all, in any order, with repeats. That generality is exactly why it cannot be a view: no single stride visits rows 2, 7 and 3 in that order.

narrow only accepts a contiguous run, and in exchange it is free. It changes the offset and the size of one axis and nothing else. x[2:6] is the same operation with nicer syntax.

The habit

When the rows you want are a range, say so. Reaching for index_select with torch.arange(start, stop) is a copy of the entire slice, and it is a surprisingly common thing to find in a data pipeline, usually because the general helper was written first and the range case went through it.

This matters most in a Dataset or a collate function, where it runs once per batch forever.

When you genuinely need index_select

Shuffled indices, sampled indices, sorted-by-something indices, repeated indices. Those are copies and there is no way around it; the data really is being rearranged. x[idx] does the same thing with the same cost.

Your task

def take_rows(x: torch.Tensor, start: int, length: int) -> torch.Tensor

Return length rows of x beginning at start, without copying.