Functional Update with .at[].set()

Replace entries and keep the original array

Ordinary JAX arrays do not support item assignment. Instead, arr.at[indices].set(values) returns an updated array value. Keep that result; the original array remains unchanged.

import jax.numpy as jnp

arr = jnp.array([1., 2., 3.])
updated = arr.at[1].set(99.)
# arr:     [1., 2., 3.]
# updated: [1., 99., 3.]

Rebinding arr to the result is also valid. It changes what the Python name refers to, not the old array value. Compiler buffer reuse is an implementation detail and does not let you observe mutation of the original ordinary array.

Pair each position with a replacement

For arr = [10, 20, 30, 40], indices = [0, 3] and values = [7, 8], set position 0 to 7 and position 3 to 8. The result is [7, 20, 30, 8]. This replaces entries rather than adding contributions.

The exercise supplies whole-number indices as float arrays. Convert them to int32 before indexing. Casting is an interface conversion, not validation: a fractional number would be truncated and would not identify the intended position reliably.

Conflicting replacements need a defined policy

If several replacements target the same position, .set does not guarantee that the last input value wins on every backend. Do not infer a portable rule from the result on one machine. This exercise uses distinct indices so each destination has exactly one replacement. Use .add when the intended operation is to accumulate contributions; that changes the operation, not just its implementation.

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.