Medium Framework

Scatter Add

Scatter-add sends source values to output positions and sums values that share a destination. Start from zeros, so output positions that receive no values stay zero.

Your task

Implement scatter_add(src, indices, n).

src is a floating-point vector of shape (k,).
indices has shape (k,) on the same device. Its values are integer-valued floats in [0, n); cast them to integers before indexing.
n is a nonnegative integer giving the output length. If n=0, k is also zero.
Return a vector of shape (n,) with src's dtype and device. Leave the inputs unchanged.

For each source position i, add src[i] to output position indices[i]:

output[j] = sum(src[i] for every i whose indices[i] == j)

For src=[2, -3, 5], indices=[1, 1, 3] and n=5, return [0, -1, 0, 5, 0]. The two contributions to position 1 add to -1. Repeated indices mean addition, not last-write-wins assignment.

The output length comes from n, not from the source length or largest index. An empty source returns n zeros; if n is also zero, return shape (0,).

Implementation

Create the zero vector with the source dtype and device. Then use target.scatter_add(0, indices.long(), src) in PyTorch or target.at[indices.astype(jnp.int32)].add(src) in JAX. A sequential loop that adds each contribution also expresses the required operation.

Do not use ordinary indexed assignment: duplicate destinations must retain all contributions. Floating-point sums may have small rounding differences when a backend adds contributions in a different order.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

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.

Detecting runtime…