JAX is a library for differentiable computation built around stateless functions. It provides low-level array operations and automatic differentiation, and it can serve as the backend for higher-level Keras models. Its stateless approach helps computation compile and scale across devices, but it also changes how you write a training loop: pass state in, then return updated state. section 3.5 section 3.2
Train a linear classifier
Start with a model that computes `jnp.matmul(inputs, W) + b` and a mean-squared-error loss. Put `W` and `b` in a state tuple passed to `compute_loss`. Applying `jax.value_and_grad(compute_loss)` produces both the loss and gradients for that state. A training step subtracts a learning-rate-scaled gradient from each weight and returns the loss alongside the new `W` and `b`. Repeated calls form the training loop; no function silently retains the weights between calls. section 3.5.4
Three practical rules make the pattern easier to follow:
- Use `jax.numpy`, commonly imported as `jnp`, for NumPy-style array operations. Its API has differences around random number generation and array assignment. section 3.5.2
- Keep changing values explicit. If a learning rate must change during training, pass it through as state too. section 3.5.4
- Use `jax.jit` to compile a training step, while allowing for harder debugging than with purely eager execution. section 3.5.4 section 3.5.5
learning_rate = 0.1 @jax.jit #1 def training_step(inputs, targets, W, b): loss, grads = grad_fn((W, b), inputs, targets) #2 grad_wrt_W, grad_wrt_b = grads W = W - grad_wrt_W * learning_rate #3 b = b - grad_wrt_b * learning_rate #3 return loss, W, b #4
The compiled training step shows the central pattern in code: obtain loss and gradients, calculate new weights, and return them instead of mutating persistent model state. section 3.5.4
The same pattern extends to a custom Keras training step on a JAX backend. A model’s `stateless_call()` returns outputs and updated non-trainable variables; `jax.value_and_grad(..., has_aux=True)` obtains gradients while retaining those updates; and an optimizer’s `stateless_apply()` returns updated trainable and optimizer variables. That is more state to thread through a function than in the TensorFlow or PyTorch versions. section 7.4.2
JAX gradients also support model interpretation: `jax.grad()` supplies image gradients for gradient ascent and gradients of a predicted class with respect to a convolutional layer’s output. JAX is a strong fit when speed and TPU-scale training matter; for a simple low-level loop, weigh those benefits against more verbose code and compilation-related debugging friction. section 10.2.3 section 10.3.3 section 3.5.5
