JAX makes program state explicit. Functions receive model parameters, optimizer state, gradients, and random-number-generator state as arguments instead of hiding them inside objects. This reduces side effects and makes computations easier to compose, transform, parallelize, cache, and debug. section 1.1.2 section 2.9
A pure function depends only on its inputs and returns the same output for the same inputs. JAX transformations such as automatic differentiation, JIT compilation, and vectorization work best when functions have no side effects. “Stateless” therefore means externally managed state, not that state disappears. section 2.9 section 1.4
This philosophy changes familiar programming habits:
- Pass neural-network parameters directly into computation functions rather than storing them as hidden object state. section 2.9
- Pass random-number-generator state explicitly; identical inputs and state produce identical “random” results. section 2.9
- Replace in-place array mutation with operations that return a modified copy, such as x = x.at[idx].set(y). section 3.3.1 section 1.2.1
- Treat PyTorch and TensorFlow’s object-oriented organization as a contrast: JAX favors functions and composable transformations, while higher-level JAX libraries can still offer familiar abstractions. section 1.3
Functional array update
NumPy permits a direct mutation such as a_np[5] = 100. JAX rejects a_jnp[5] = 100 because arrays are immutable. Write a new binding instead: a_jnp = a_jnp.at[5].set(100). The original array remains unchanged, while JIT compilation can still apply low-level updates efficiently when the original copy is no longer needed. section 3.3.1
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: ...
Shows the NumPy mutation, JAX error, and functional replacement side by side. section 3.3.1
