Zero the padding, keep the tensor

The previous problem showed advanced indexing failing to write through. This is the other half: basic slicing writes through whether you wanted it to or not.

x = torch.ones(2, 5)
region = x[:, 3:]
region.zero_()
x                     # the last two columns of x are now zero

region was never a separate tensor. It was x with a different shape and offset.

Both directions are bugs, in different code

Wanting the write and not getting it is the advanced-indexing trap. Not wanting the write and getting it is this one, and it usually shows up as a function that “helpfully” normalises its input and corrupts the caller’s data:

def normalise(x):
    x[:, 0] -= x[:, 0].mean()    # the caller's tensor just changed
    return x

The caller passed a tensor and got their tensor back, modified. If that tensor was a training batch someone else also holds, you now have a bug that reproduces once every few epochs.

The rule

Decide whether you are mutating or deriving, and say so:

out = x.clone()      # deriving. Now nothing you do escapes.
x[...] = ...         # mutating. Deliberate, documented, in place.

clone() copies the buffer. detach() does not, and is about autograd rather than memory, which is a distinction the autograd track returns to.

Your task

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

Return a new tensor equal to x with every column from index keep onwards set to zero. x itself must come back unchanged.

The starter slices and zeroes, which is the right idea applied to the wrong tensor.