JAX arrays cannot be modified in place. This follows functional programming: instead of changing an existing value, create a new array from it. NumPy permits updates such as arr[i] += 10, but JAX requires a functional alternative. section 3.3.1 section 1.2.1
Replacing an indexed update
a_jnp = jnp.array(range(10)) a_jnp = a_jnp.at[5].set(100) # Array(100, dtype=int32) The original array remains conceptually unchanged; the expression returns the modified copy. JAX's error message for a_jnp[5] = 100 points to this .at[] syntax.
Listing 3.21 Updating an element of a tensor
a_jnp = jnp.array(range(10)) a_np = np.array(range(10)) a_np[5], a_jnp[5] >>> (5, Array(5, dtype=int32)) a_np[5] = 100 #1 a_np[5] >>> 100 try: a_jnp[5] = 100 #2 except TypeError as e: print(e) >>> '<class 'jaxlib.xla_extension.ArrayImpl'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: ...
Compare successful NumPy mutation with JAX's informative item-assignment error and its recommended functional replacement. section 3.3.1
