Skip to content

← Indexing and Selection step 3 of 9

Medium Primitives

When two writes land on one slot

Scattering is gathering run backwards: instead of reading from positions an index names, you write to them. The interesting question is what happens when the index names the same position twice.

index   [1, 1, 3]
source  [10, 20, 30]
into    [0, 0, 0, 0, 0]
out.scatter(0, index, source)       # [0, 20,  0, 30, 0]
out.scatter_add(0, index, source)   # [0, 30,  0, 30, 0]

scatter writes. Two writes to slot 1, and one of them wins. Which one is not specified when the duplicates are in the same call, so the 20 you see here is not a promise.

scatter_add accumulates. Slot 1 gets 10 + 20, deterministically.

Why this is the whole problem

Counting, binning and segment-summing are all “several sources, one destination”, which means duplicate indices are not an edge case, they are the point:

counts = torch.zeros(n_classes)
counts.scatter_add_(0, labels, torch.ones_like(labels, dtype=torch.float))

Write scatter_ there and you have built a presence check that reports 1 for every class that occurred at all. It will look plausible, it will train, and the numbers will be wrong.

This is the same distinction as index_put_ with and without accumulate=True, and as bincount versus unique. Whenever you write through an index, ask what you want to happen on a collision, because something is going to happen either way.

Your task

def bin_counts(labels: list[int], n_bins: int) -> torch.Tensor

Return a float tensor of length n_bins where entry i is how many times i appears in labels.

The starter scatters ones. Every bin that occurred gets a 1, however often it occurred.