Skip to content

← Indexing and Selection step 7 of 9

Easy Primitives

Rows 2 through 5, cheaply

Jump to the task

Selecting a consecutive range of rows can avoid copying their data. The operation you choose determines whether the result shares storage.

import torch

x = torch.arange(12, dtype=torch.float32).reshape(6, 2)
view = x.narrow(0, start=2, length=2)
copied = x.index_select(0, torch.tensor([2, 3]))
print(view.tolist())    # [[4.0, 5.0], [6.0, 7.0]]
print(copied.tolist())  # [[4.0, 5.0], [6.0, 7.0]]

Both results contain the same values, but narrow returns a view and index_select returns a copy. Even consecutive indices do not make index_select share storage.

What sharing means

For this valid range, x[2:4] also returns a view. A view describes the selected region using shape, strides and a storage offset; it does not copy the elements. Creating that view still has bookkeeping overhead.

view.add_(100)
print(x[2:4].tolist())   # [[104.0, 105.0], [106.0, 107.0]]
print(copied.tolist())  # [[4.0, 5.0], [6.0, 7.0]]

A view is useful when sharing is intended. A copy is useful when later changes should be independent. Keeping a small view alive can also keep the original tensor's larger storage alive.

Choose the operation that expresses the selection

narrow selects a consecutive interval along an axis. Those rows need not be contiguous in memory for a strided tensor; the operation preserves the relevant strides and shares storage.

index_select accepts an explicit index vector, including reordered and repeated positions. Use it when that is the selection you need. Some regular patterns can instead be expressed as slices, so a list of indices is not itself proof that copying is unavoidable.

Your task

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

Return a view containing length rows of the two-dimensional tensor x, starting at row start. Inputs specify a valid nonnegative start and positive length within the row bounds. Preserve the columns and share storage with x. The starter returns the right values but copies them.