Immutable Arrays: Functional Updates with .at

Easyjax

JAX arrays are immutable — x[i] = v raises an error. Updates are expressed functionally: x.at[i].set(v) returns a new array with the change applied, leaving x untouched.

Implement replace_at(x, indices, values):

x: a 1-D jnp array
indices: a list/array of integer positions
values: the new values for those positions (same length as indices)

Return the updated array without modifying x. The test prints both the result and the original array afterward, so a mutating implementation fails even if the returned value looks right.

Example:

Input:

x = jnp.array([10.0, 20.0, 30.0, 40.0]), indices = [1, 3], values = [-1.0, -2.0]

Output:

[10.0, -1.0, 30.0, -2.0]
[10.0, 20.0, 30.0, 40.0]

Reasoning:

Positions 1 and 3 receive the new values; the original array is unchanged because .at[].set() is out-of-place.

Learn About topic

Contributors:

Deep-ML

Contribute