The speed-and-scale backend
The JAX backend (JAX 0.4.20+) is frequently the fastest option for GPU and TPU training, because JAX compiles your computation through XLA into one fused, hardware-tuned kernel instead of dispatching op-by-op. It's the engine behind much of Google's large-scale training, and it carries three things into Keras that matter once your model gets serious:
- JIT compilation via XLA for maximum throughput
- Functional transforms —
jax.vmap,jax.grad,jax.pmap— composable on your own math - The strongest support for
keras.distribution(multi-GPU/TPU data and model parallelism)
The functional-purity tax
JAX's power comes from a constraint: transformed functions must be pure — no hidden mutable state. Keras model objects do hold state (weights, optimizer slots, metric tallies), so to run them inside a JIT/grad transform you pass that state in and get the updated state back, rather than mutating in place. That's the stateless_call pattern in the Code block. If you only ever call fit() you may never touch it; the moment you write a custom JAX training step, it becomes the core idiom.