When two writes land on one slot

When several values target the same destination, decide whether to replace the stored value or add the contributions. Counting requires addition.

import torch

index = torch.tensor([1, 1, 3], dtype=torch.long)
source = torch.tensor([10.0, 20.0, 30.0])
out = torch.zeros(5)
out.scatter_add_(0, index, source)
print(out.tolist())  # [0.0, 30.0, 0.0, 30.0, 0.0]

Position 1 receives both 10 and 20. Position 3 receives 30. scatter_ instead writes source values, and with duplicate indices its choice of surviving value is unspecified. Do not rely on “the last write wins.”

Build a histogram

Give each label a contribution of one and add it to that label's bin:

labels = torch.tensor([1, 1, 3], dtype=torch.long)
counts = torch.zeros(5)
counts.scatter_add_(0, labels, torch.ones(3))
print(counts.tolist())  # [0.0, 2.0, 0.0, 1.0, 0.0]

Replacing rather than adding these identical ones would leave every occupied bin at one, losing the number of occurrences. Repeated labels are central to a histogram, so the collision behavior matters.

scatter_add_ specifies addition, but does not guarantee identical floating-point results on every device. CUDA accumulation can be nondeterministic because contributions may be combined in a different order. The small integer counts in this CPU exercise are exactly representable.

torch.bincount is another option for nonnegative integer labels. It produces integer counts for unweighted input; this exercise asks for a floating-point output and illustrates accumulation through an index.

Select values using coordinates

A slice selects a regularly spaced region of an array. A gather selects entries using supplied indices. Check whether those indices select whole rows or columns, or a different entry within each row.

For [[10, 11], [20, 21]], choosing column 1 from the first row and column 0 from the second gives [11, 20]. Choosing columns [1, 0] from every row instead gives [[11, 10], [21, 20]]. These operations use the same index values but have different meanings and output shapes.

Scatter writes values to selected destinations

Scatter takes values and destination indices. The operation may replace existing values or combine them using a reduction such as addition.

Starting with [0, 0, 0], scatter-add of values [2, 3] at indices [1, 1] gives [0, 5, 0]. Both values contribute to the repeated destination. Repeated assignment is not equivalent to this sum; follow the operation's rule for repeated indices.

Selecting with a mask differs from replacing with a mask

For values [10, 20, 30] and mask [True, False, True], selecting masked entries gives [10, 30]. Replacing masked entries with zero gives [0, 20, 0]. Selection can change the shape; replacement preserves it. A selection whose size depends on runtime values needs particular care in compiled code with fixed output shapes.

Check whether an update changes the original array

JAX indexed updates return a new array value. In PyTorch, basic slicing can return a view that shares storage, while advanced indexing returns a copy. Modifying a selected copy does not modify the source. Direct indexed assignment is a separate operation, so distinguish changing a selection from assigning into the original.