JAX and the Google stack

Models and offerings #

Who they are

This hub is for people choosing between JAX and PyTorch for new work, and for anyone who has been handed a TPU allocation. JAX started inside Google Research in 2018 as autograd plus XLA: NumPy functions that could be differentiated, vectorised and compiled. Google DeepMind adopted it for most of its research, and the Gemini and Gemma models described on the Google hub are trained with it. TensorFlow, Google’s earlier framework, is in maintenance for new research and most of its users have moved to JAX, PyTorch or Keras 3.

JAX’s share of the field is small next to PyTorch but concentrated: TPU users, reinforcement learning, scientific computing and anyone who likes pure functions. The ecosystem around it is modular rather than batteries-included, which is both its appeal and its tax.

Tooling and APIs

pip install jax gives you the CPU build. GPU builds are jax[cuda13] (or the CUDA 12 variant) on Linux x86_64 and aarch64, and jax[tpu] on Cloud TPU VMs. The support table in the installation guide is honest and worth reading: Linux CPU and NVIDIA GPU are fully supported, macOS is CPU only (the Apple GPU path is listed as experimental and the same page says Mac GPU is not supported), Windows is CPU only with GPU through WSL2 marked experimental, AMD ROCm works on Linux x86_64, and Intel GPUs are experimental. JAX 0.12.0 (19 October 2026) requires NumPy 2.2 or newer and an x86-64 CPU of Haswell vintage or later. Python 3.11 support was dropped in 0.11.0, so you need 3.12 or newer.

JAX itself has no nn.Module. Flax supplies that, and since 2024 the recommended API is NNX, which holds state on Python objects the way PyTorch does. The older Linen API is purely functional and still maintained. Optax supplies optimisers, Orbax checkpoints, and Chex and Grain cover testing and data loading. Keras 3 sits above all of this if you want model.fit. The whole stack is Apache-2.0.

The core idea in a few lines:

# snippet
import jax
import jax.numpy as jnp

def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)

grad_loss = jax.jit(jax.grad(loss))          # compile the gradient function
batched = jax.vmap(loss, in_axes=(None, 0, 0))  # per-example losses, no loop
w = jnp.zeros(4)
x = jax.random.normal(jax.random.key(0), (256, 4))
y = x @ jnp.array([1.0, -2.0, 0.5, 3.0])
for _ in range(200):
    w = w - 0.1 * grad_loss(w, x, y)
print(w)   # approaches [1, -2, 0.5, 3]

JAX is not installed in the Field Guide pipeline environment, so the block above is a syntax-checked snippet. The pain points are the flip side of the design: every function you jit must be pure, in-place updates become x.at[i].set(v), random numbers need explicit keys, and the first call of a compiled function can take seconds or minutes. Error messages from inside XLA are worse than PyTorch’s.

TPU pricing is per chip-hour and varies by generation (v5e, v5p, Trillium v6e, Ironwood v7) and region. We could not read Google’s pricing table at the time of writing, so see Google Cloud’s TPU pricing page. Hugging Face Inference Endpoints also rent TPU v5e slices by the hour.

Migration notes

From PyTorch: the big shift is that parameters are data you pass in, not attributes you mutate. A PyTorch module becomes a Flax NNX module (close to home) or a pure function plus a parameter pytree (further away). optimizer.step() becomes tx.update(grads, state) followed by optax.apply_updates(params, updates). Data loaders are not provided, so most people keep torch.utils.data or use Grain. Our frameworks reference has a side-by-side table.

Model weights move between the two through safetensors or NumPy arrays, since both are plain tensors. Hugging Face Transformers v5 dropped its Flax model classes, so a JAX port of a Hub model now comes from a separate project (MaxText, EasyDeL or the model’s own repository) rather than FlaxAutoModel. Converting a JAX checkpoint to PyTorch means mapping parameter names and transposing linear layers, which is tedious but mechanical.

To TPU from GPU: jax.sharding and shard_map replace device-specific collectives, and the same code runs on one chip or a pod slice. PyTorch users can reach TPUs through PyTorch/XLA, which works but trails native JAX on features.

Watch for

The API moves. The 0.11 and 0.12 release notes remove jax.custom_remat, jax.experimental.shard_alike and the cloud_tpu_init module, change return types of meshgrid and broadcast_arrays to tuples, and make jnp.empty return genuinely uninitialised memory. Pin jax, jaxlib and flax together.

TPUs are a Google Cloud product and nowhere else. Code written with JAX is portable, but checkpoints written by Orbax in TPU-sharded layouts need care to reload on GPUs. Free TPU access through Colab and the TPU Research Cloud exists but is capacity-limited and comes with usage terms.

Flax has two APIs and the documentation covers both, so check which one a tutorial uses before copying it. Keras 3 on the JAX backend is a clean way in for Keras users but adds a layer whose bugs are harder to trace. There is no telemetry in the libraries. Colab and Cloud TPU VMs are Google services with their own terms.