Functional Update with .at[].add()

Add contributions to selected positions

Scatter-add collects values by destination. Several contributions can target the same entry, as in a histogram or an embedding gradient. arr.at[indices].add(values) adds them to arr's existing contents and returns the updated array value.

import jax.numpy as jnp

arr = jnp.array([10., 20., 30.])
indices = jnp.array([0, 0, 2])
values = jnp.array([1., 3., -5.])
updated = arr.at[indices].add(values)
# [14., 20., 25.]
# index 0: 10 + 1 + 3; index 2: 30 - 5

Arr stays [10, 20, 30]. Keep or return updated; discarding the result does not change arr. Starting from nonzero values matters: this is accumulation into the original array, not a sum into a new zero array.

Repeated indices accumulate

All contributions to a destination participate in the sum. This differs from .set, which replaces values and has no portable last-write rule for conflicting destinations.

Accumulation order can vary with backend and execution. Floating-point addition is not associative, so repeated-index sums are not guaranteed bitwise identical across hardware. The mathematical sum includes every update; its floating-point rounding can depend on their order.

Indexing is not arbitrary wraparound

The default .at update behavior does not wrap a large positive index around the array. Out-of-range updates are normally dropped under the default promise_in_bounds mode, which does not validate the promise. Negative indices count from the end by default, as in ordinary Python indexing. These are separate behaviors. Explicit drop or clip modes are available when a task needs a particular boundary policy.

Here every index is nonnegative and in range, so boundary handling is not part of the exercise. Convert the supplied whole-number float indices to int32; casting fractional values would not validate them.

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.