Dynamic Update Slice

Replace a contiguous block

lax.dynamic_update_slice(x, update, starts) returns x with a block replaced by update. The start position can change at runtime under JIT. Values outside the block stay the same, and the result keeps x's shape. This is replacement, not addition or insertion that grows the array.

import jax
import jax.numpy as jnp
from jax import lax

@jax.jit
def replace_two(x, start, patch):
    return lax.dynamic_update_slice(x, patch, (start,))

x = jnp.array([1., 2., 3., 4., 5., 6.])
patch = jnp.array([7., 8.])
replace_two(x, 2, patch)  # [1., 2., 7., 8., 5., 6.]
replace_two(x, 0, patch)  # [7., 8., 3., 4., 5., 6.]
# x still contains [1., 2., 3., 4., 5., 6.]

Capture or return the result. JAX arrays do not support assignment such as x[2:4] = patch, even with constant indices. The functional spelling x.at[2:4].set(patch) works for a static slice, but a slice whose bounds are traced needs a different approach, such as dynamic_update_slice.

The block must fit, but the requested start can move

For a vector of length n and an update of length m, require m <= n. If the requested start is too far right, JAX shifts it left to n - m so the entire update fits. This is documented behavior, not undefined behavior. In the example above, start 5 produces [1, 2, 3, 4, 7, 8]. The patch is neither cut short nor wrapped to the beginning.

By default, negative starts count from the end before this adjustment. Start -1 gives the same result as start 5; start -3 replaces positions 3 and 4, giving [1, 2, 3, 7, 8, 6]. A start below -n is adjusted to zero. An update larger than x cannot be made to fit by shifting it.

Shapes and values have different roles

Pass one integer scalar start per axis. This exercise converts the supplied start to int32, but the API does not require every start to use that particular integer dtype. The update must match x's rank and dtype; it is not broadcast from a scalar or a shorter dimension.

Under ordinary JIT, x and update have fixed shapes for a compiled specialization, while their values and the start can vary. There is no separate static size argument: the block size comes from update.shape.

Dynamic update is not an inverse that can recover overwritten values. It replaces them. Immutability describes the visible array behavior, not a guarantee that the compiler allocates a separate buffer every time.

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.