The Un-Updatable Array

In-place mutation like x[0] = 5 raises an error on a JAX array. What is the primary reason JAX arrays are immutable, requiring functional updates like x = x.at[0].set(5) instead?

It prevents accidental type coercion between float32 and float64

Immutability preserves the referential transparency that tracing and compiling functions safely depends on

Immutable arrays use less GPU memory than mutable NumPy arrays

XLA-targeted hardware backends are physically incapable of overwriting memory buffers