2 / 8 JAX & XLA Mental Model

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?

A

It prevents accidental type coercion between float32 and float64

B

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

C

Immutable arrays use less GPU memory than mutable NumPy arrays

D

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