JAX is a low-level library for numerical computation and automatic differentiation. It can provide the tensor operations and gradients beneath Keras, or support a custom training loop directly. Its defining habit is to keep functions stateless: pass model weights in, then return updated weights rather than changing hidden state. section 3.2 section 3.5 section 3.5.4
Train a small linear model
This R example fits four input–target pairs. The loss function receives the weights as explicit state; `value_and_grad()` returns its loss and gradients. The compiled training step returns new weights for the next iteration. ```r library(reticulate) library(keras3) use_backend("jax") jax <- import("jax") jnp <- jax$numpy inputs <- jnp$array(matrix(c(0, 0, 0, 1, 1, 0, 1, 1), ncol = 2, byrow = TRUE), dtype = "float32") targets <- jnp$array(matrix(c(0, 1, 1, 2), ncol = 1), dtype = "float32") compute_loss <- function(state, inputs, targets) { .[W, b] <- state predictions <- jnp$matmul(inputs, W) + b jnp$mean(jnp$square(targets - predictions)) } grad_fn <- jax$value_and_grad(compute_loss) training_step <- jax$jit(\(inputs, targets, W, b) { .[loss, grads] <- grad_fn(list(W, b), inputs, targets) .[grad_wrt_W, grad_wrt_b] <- grads tuple(loss, W - 0.1 * grad_wrt_W, b - 0.1 * grad_wrt_b) }) W <- jnp$array(matrix(runif(2), ncol = 1), dtype = "float32") b <- jnp$array(c(0), dtype = "float32") for (step in seq_len(40)) { .[loss, W, b] <- training_step(inputs, targets, W, b) } as.array(loss) ``` section 3.5.1 section 3.5.4
Three details shape everyday JAX code:
- Random operations take an explicit PRNG key. Reusing a key reproduces the same result; split it when a new key is needed. section 3.5.3
- Arrays are treated without in-place mutation in JAX functions. Pass state as arguments and return its new value. section 3.5.3 section 3.5.4
- `jax$grad()` computes gradients; `jax$value_and_grad()` supplies both a value and its gradients. `jax$jit()` compiles a function for faster execution. section 3.5.3
With Keras on the JAX backend, a custom training step must also thread state through the computation: a stateless model call handles trainable and non-trainable variables, and a stateless optimizer call returns updated optimizer and model variables. JAX gradients can also target a predicted class when interpreting a convolutional network. section 7.4.2 section 10.3.3
JAX is used for large-scale training, including Keras data parallelism across GPUs. Compilation and stateless execution support speed and scaling, but compilation can complicate debugging, and low-level loops can be verbose. Multi-GPU speedups also depend on having a sufficiently large batch. section 3.5.5 section 18.2.2
keras$distribution$set_distribution( keras$distribution$DataParallel() )
The one-line Keras data-parallel configuration for the JAX backend. section 18.2.2
