Skip to content

← Indexing and Selection step 4 of 9

Easy Primitives

Zero the padding, keep the tensor

Jump to the task

Basic slicing creates a view that shares the source tensor's storage. Changing the view can therefore change the source:

import torch

x = torch.ones(2, 5)
region = x[:, 3:]
region.zero_()
print(x.tolist())  # [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0, 0.0]]

region is a separate tensor object, but its values live in the same storage as x. Its shape, strides and storage offset describe the selected part of that storage.

Choose whether the caller's input should change

In-place operations are useful when mutation is intended. If a function promises to preserve its input, first create independent data storage:

x = torch.tensor([[1., 2., 3.]])
out = x.clone()
out[:, 1:] = 0
print(out.tolist())  # [[1.0, 0.0, 0.0]]
print(x.tolist())  # [[1.0, 2.0, 3.0]]

clone copies the data. detach shares storage and therefore does not protect the input from mutation. Cloning can still preserve a gradient connection to the original tensor; storage independence and gradient tracking are separate concerns.

Reading a selection and assigning to it differ

Basic slices share storage when read. Advanced indexing, such as a list of column indices, normally returns a copy when read. However, direct indexed assignment like x[:, [1, 3]] = 0 writes into x. Do not assume that every use of advanced indexing prevents mutation.

Likewise, an assignment to x[:, keep:] modifies x directly. The starter uses exactly that form, so it returns the desired values by changing the caller's input.

Your task

def zero_tail(x: torch.Tensor, keep: int) -> torch.Tensor

Given a nonempty strided 2-D tensor and 0 <= keep <= x.size(1), return a new tensor with columns from keep onward set to zero. Preserve the values before keep and leave x unchanged. The result must have independent data storage, including when keep equals the width and no values change.

The harness checks the output values, the input after the call and storage independence. Copy first, then modify the copy.