JAX isn't an island — it sits at the center of a growing ecosystem of libraries that build higher-level functionality on top of JAX's core transformations.
Key Libraries in the JAX Ecosystem
- Flax (by Google) — Neural network library. The newer Flax NNX API provides a more Pythonic, mutable-style interface for defining models while maintaining JAX's functional core.
- Optax — Gradient processing and optimization library. Provides Adam, SGD, learning rate schedules, gradient clipping, and more as composable transformations.
- Orbax — Checkpointing and model export utilities.
- Chex — Testing utilities for JAX code (shape/dtype assertions, test variants).
- Equinox (by Patrick Kidger) — A lightweight neural network library that represents models as PyTrees. Popular in the research community for its elegance.
- Diffrax — Differential equation solvers built on JAX.
- BlackJAX — Bayesian inference and MCMC sampling.
- Mctx — Monte Carlo tree search (used in game-playing AI).
Who Uses JAX?
JAX is the primary framework at Google DeepMind, which means it's behind some of the most high-profile AI results in recent years. Beyond Google, it's widely used at research labs worldwide for work that demands high performance and mathematical flexibility. Companies and teams that need to train on TPU pods often choose JAX because of its first-class TPU support.
Installation
For CPU-only (great for learning):
# CPU only
pip install jax
# GPU (NVIDIA, CUDA 12)
pip install jax[cuda12]
# TPU (Google Cloud)
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Verify your installation:
import jax
print(jax.__version__) # e.g., 0.5.0 or later
print(jax.devices()) # [CpuDevice(id=0)] or [CudaDevice(id=0)]
print(jax.default_backend()) # 'cpu', 'gpu', or 'tpu'
# Quick test
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
print(jax.grad(lambda x: jnp.sum(x ** 2))(x)) # [2. 4. 6.]
Versioning
Since version 0.5.0 (January 2025), JAX uses effort-based versioning (EffVer). In this scheme, the three-number version MACRO.MESO.MICRO signals how much effort you'll need to adapt to changes:
- Micro bump (e.g., 0.5.0 → 0.5.1): little to no effort needed
- Meso bump (e.g., 0.5.x → 0.6.0): some small effort to update code
- Macro bump (e.g., 0.x.y → 1.0.0): significant effort may be required
As of early 2026, the latest JAX release is in the 0.9.x series. JAX is still in its zero-version phase (0.x.y), where the meso number (the second digit) acts like a macro version. When JAX reaches a stability plateau at 0.9.x for several months, the team plans to release 1.0.0.
💡 Why This Matters
JAX's versioning tells you that the API is still evolving. When following tutorials or Stack Overflow answers, always check what JAX version they were written for. APIs like jax.tree_util have been simplified to jax.tree, and sharding APIs have been modernized significantly in recent versions. This tutorial targets the modern API (0.5+).