Skip to content

← Indexing and Selection step 1 of 9

Medium Primitives

Update selected rows, including repeated indices

Jump to the task

Reading a slice and reading a list of indices have different storage behavior. A basic slice is a view, so mutating it can change the original:

import torch

a = torch.zeros(5)
view = a[1:3]
view.add_(1)
print(a.tolist())  # [0.0, 1.0, 1.0, 0.0, 0.0]

Advanced indexing with a list of integer indices returns a copy. Updating that copy leaves the source unchanged:

b = torch.zeros(5)
selected = b[[1, 2]]
selected.add_(1)
print(selected.tolist())  # [1.0, 1.0]
print(b.tolist())         # [0.0, 0.0, 0.0, 0.0, 0.0]

Advanced reads copy even when the chosen positions happen to be consecutive. The indexing operation's semantics determine sharing, not just whether those positions could have been described by a slice.

Writing directly is different

With distinct indices, an indexed assignment writes back to the original:

b[[1, 2]] += 1
print(b.tolist())  # [0.0, 1.0, 1.0, 0.0, 0.0]

This augmented assignment includes a write-back step. In contrast, selected.add_(1) only writes into the previously created copy.

Repeated indices need extra care. x[rows] += amount is not a reliable way to add once per occurrence: advanced assignment does not accumulate duplicate destinations by default. Use an accumulation operation such as index_add_ when duplicates should contribute multiple times.

Your task

def add_to_rows(x: torch.Tensor, rows: list[int], amount: float) -> torch.Tensor

x is a two-dimensional CPU float32 tensor without gradient tracking. rows contains valid nonnegative row indices. Add amount to every element of each listed row, in place, and return x.

A row mentioned twice must receive two additions. For example, starting from zeros with rows=[1, 1] and amount=5.0 leaves every element of row 1 equal to 10.0. Rows not listed must remain unchanged.