TopolearnJAX arrays cannot be mutated, so every index update in ported NumPy code needs this pattern.
NumPy code often does x[0] = 10 or x[mask] += 1. In JAX that raises TypeError because jax.Array is immutable. Immutability is what lets JAX treat functions as pure mathematical mappings, which in turn makes grad, jit and vmap possible.
Instead use the indexed update property .at:
x.at[i].set(v)returns a copy ofxwith elementireplaced.x.at[i].add(v),.multiply(v),.min(v),.max(v)combine with the existing value.x.at[i].get(mode="fill", fill_value=...)reads with explicit out-of-bounds behaviour.
The index can be anything NumPy indexing accepts: slices, integer arrays, boolean masks of static shape.
Two important semantics:
- With repeated indices,
.addaccumulates (x.at[[0, 0]].add(1)adds 2), unlike NumPy'sx[[0, 0]] += 1, which adds only once. - Updates at out-of-bounds indices are silently dropped by default.
"Returns a copy" sounds slow, but inside a jit-compiled function XLA sees that the old array is no longer used and performs the update in place. Outside jit each update really does allocate, so batch your updates or compile the function.
For element-wise conditional replacement, jnp.where(cond, a, b) is usually clearer than masked .set.
Example
import jax.numpy as jnp
x = jnp.zeros(5)
try:
x[0] = 1.0
except TypeError:
print("TypeError") # JAX arrays are immutable
y = x.at[0].set(1.0)
print(x) # [0. 0. 0. 0. 0.] original untouched
print(y) # [1. 0. 0. 0. 0.]
z = y.at[1:3].add(5.0)
print(z) # [1. 5. 5. 0. 0.]
print(z.at[jnp.array([4, 4])].add(1.0)) # [1. 5. 5. 0. 2.] duplicates accumulate
print(z.at[10].set(9.0)) # [1. 5. 5. 0. 0.] out-of-bounds update dropped
print(jnp.where(z > 2, 0.0, z)) # [1. 0. 0. 0. 0.]Remember
x.at[idx].set(v)returns a new array; the original is unchanged.addwith repeated indices accumulates every occurrence- Inside
jit, XLA turns these copies into in-place updates
Watch out for
- Expecting an out-of-bounds
.setto raise an IndexError; it is silently dropped.