jax review
JAX turns NumPy-shaped Python functions into programs that can be differentiated, vectorized, compiled with XLA, and sharded across accelerators. The central tools are `grad`, `vmap`, `jit`, and explicit array sharding; they compose around pure functions rather than a model class. Version 0.11.1 adds `jax.numpy.top_k`, rejects exported artifacts outside its compatibility window, and changes several NumPy-alignment details. Our CPU install worked, though the six-package environment occupied 525 MB before any GPU stack was added.
JAX 0.11.1 installed in 2 seconds and imported in 2.19 seconds on our CPU box, yet its 6-package environment consumed 525 MB before accelerator extras. Choose it for transformable numerical programs and TPU or multi-device work; choose PyTorch when existing models and serving integrations matter more.
We installed it
| Install | ✓ · 2s | 6 packages on disk · 525 MB |
| Import | ✓ | import jax in 2.19s · pure Python · py.typed · requires Python >=3.12 |
| Known vulns | 0 | (pip-audit) |
Answers from our run
Does jax install cleanly?
Yes. In a fresh container with an empty cache, pip install jax finished in 2 seconds, leaving 6 packages and 525 MB on disk. pip-audit reported no known vulnerabilities.
What does jax need to run?
Python >=3.12, and nothing compiled: it is pure Python. In our run import jax succeeded in 2.19s, and the package ships py.typed for type checkers.
jax or torch: which should you use?
torch: Use it when pretrained models, eager debugging, and common production serving stacks are the priority. JAX 0.11.1 installed in 2 seconds and imported in 2.19 seconds on our CPU box, yet its 6-package environment consumed 525 MB before accelerator extras.
When should you not use jax?
Your work mainly consumes pretrained models and deployment recipes. PyTorch has the larger ready-made model and serving path, so porting may dominate the project.
Use it if
- You write numerical functions whose gradients, batches, and compiled variants should come from composable transformations.
- One program must move between CPU, supported NVIDIA or AMD GPUs, and Google TPU with explicit sharding available when scale demands it.
- Scientific simulation, reinforcement learning, or original model research benefits from XLA fusion and array-first code.
- Your team accepts functional updates, explicit random keys, fixed-shape compilation, and device-aware debugging.
- Your work mainly consumes pretrained models and deployment recipes. PyTorch has the larger ready-made model and serving path, so porting may dominate the project.
- The team expects ordinary mutable NumPy semantics. JAX arrays are immutable, random state is passed as keys, and traced values cannot drive arbitrary Python control flow.
- Inputs are highly ragged or shapes change continually. `jit` specializes and compiles for new shapes and static values, which can turn flexibility into repeated compile cost.
- You need native Windows NVIDIA GPU support. The official matrix lists it as unavailable, with WSL2 support still experimental.
- A stable serialized-program format is mandatory. Version 0.11.1 refuses exports older than January 15, 2026 by default and removes two optimization-effort flags in favor of `EffortLevel`.
Setup reality
Our clean Python 3.12 Bookworm install of JAX 0.11.1 finished in 2 seconds. It left 6 packages using 525 MB, and import jax worked in 2.19 seconds. We measured 26 direct dependencies, a typed pure-Python jax package with py.typed, and 0 known vulnerabilities from pip-audit. Python 3.12 or newer is required for this release.
That measurement covers the default CPU path. Accelerators use different extras and compatibility constraints: the README gives jax[cuda13] for NVIDIA, jax[tpu] for TPU, and jax[rocm7-local] for AMD on Linux. Native Windows has CPU support but no NVIDIA GPU support; WSL2 NVIDIA is experimental. Driver, CUDA, plugin, and jaxlib compatibility must match the install guide, so pin the complete environment rather than jax alone.
A jitted function traces and compiles on its first call for a given shape and static-argument combination. Later calls may be fast, but timing an asynchronous dispatch without .block_until_ready() measures queueing instead of computation. Shape changes and new static values can compile new executables. Keep build warm-up, compilation cache policy, and representative input shapes in performance tests.
Inside transformed functions, arrays do not mutate in place, Python side effects may run only while tracing, and data-dependent branches need jax.lax control flow. Each random draw needs a fresh split key. Version 0.11.1 also enforces an export compatibility cutoff: modules serialized before January 15, 2026 fail deserialization unless the temporary override flag is set. Re-export stored programs instead of treating that flag as permanent storage policy.
Patterns
Differentiate a scalar loss compute-gradient
import jax
import jax.numpy as jnp
def loss(w, x, y):
pred = jnp.dot(x, w)
return jnp.mean((pred - y) ** 2)
g = jax.grad(loss)(w, x, y)
gw, gx = jax.grad(loss, argnums=(0, 1))(w, x, y)`grad` targets argument 0 unless `argnums` says otherwise. In 0.11.1, a non-scalar output error points you toward a reduction or `jax.jacobian`.
Compile repeated array work jit-compile
@jax.jit
def step(x):
return x * x + 2.0 * x
x = jnp.ones((5000, 5000))
step(x).block_until_ready()
step(x).block_until_ready()The first call traces and compiles this shape. Use `block_until_ready()` in benchmarks because device dispatch is asynchronous.
Map a function over a batch axis vectorize-batch
def predict(w, x):
return jnp.dot(w, x)
batched = jax.vmap(predict, in_axes=(None, 0))
out = batched(w, xs)`in_axes=(None, 0)` shares `w` and maps axis 0 of `xs`; no Python loop is created around individual examples.
Split a random key for each draw random-numbers
key = jax.random.key(42)
key, sample_key = jax.random.split(key)
x = jax.random.normal(sample_key, (3, 3))
worker_keys = jax.random.split(key, 8)Reusing the same key repeats the same random result. Split explicitly before each independent draw or worker.
Return a value and its gradient together value-and-grad
loss_value, grads = jax.value_and_grad(loss)(params, x, y)
(loss_value, metrics), grads = jax.value_and_grad(
loss_with_metrics, has_aux=True
)(params, x, y)`has_aux=True` keeps metrics outside the differentiated scalar while sharing the forward evaluation used for the gradient.
Apply an update across a PyTree update-pytree-params
grads = jax.grad(loss)(params, x, y)
params = jax.tree.map(
lambda parameter, gradient: parameter - 0.01 * gradient,
params, grads,
)Both trees must have matching structure. `jax.tree.map` is the current public spelling for mapping leaves.
Update array positions functionally functional-array-update
x = jnp.zeros((5,))
y = x.at[2].set(7.0)
z = x.at[1:3].add(1.0)
w = x.at[0].max(3.0)JAX arrays reject `x[2] = value`. Under `jit`, `.at` updates can lower to in-place device operations while Python still sees immutable values.
Declare a compile-time argument jit-static-args
from functools import partial
@partial(jax.jit, static_argnums=(1,))
def rollout(state, num_steps):
for _ in range(num_steps):
state = step(state)
return stateEach distinct static value creates another specialization. Keep the value set bounded or use a compiled control-flow primitive.
Compile a long recurrent loop scan-loop
from jax import lax
def step(carry, x):
carry = carry * 0.9 + x
return carry, carry
final, history = lax.scan(step, init=0.0, xs=data)`lax.scan` compiles one loop body instead of unrolling every iteration, which controls compile size for long sequences.
Inspect runtime values under jit debug-inside-jit
@jax.jit
def f(x):
y = x + 1
jax.debug.print('y = {y}', y=y)
return y * 2
with jax.disable_jit():
f(3.0)Plain `print` observes tracing and often displays tracer objects. `jax.debug.print` runs with compiled values; `disable_jit` helps isolate Python logic.
Place a batch across a device mesh shard-across-devices
from jax.sharding import PartitionSpec as P, NamedSharding
mesh = jax.make_mesh((len(jax.devices()),), ('data',))
sharding = NamedSharding(mesh, P('data'))
batch = jax.device_put(batch, sharding)
out = jax.jit(train_step)(params, batch)The device count and topology determine whether this one-axis mesh fits. Inspect arrays with `jax.typeof` before assuming the compiler chose the intended layout.
Use the NumPy-aligned top_k added in 0.11.1 find-top-values
values, indices = jax.numpy.top_k(scores, 5, axis=-1)
print(values.shape, indices.shape)`jax.numpy.top_k` arrived in 0.11.1 to match NumPy 2.6. Pin the JAX floor if shared code calls it.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| torch | PyPI | Use it when pretrained models, eager debugging, and common production serving stacks are the priority. |
| tensorflow | PyPI | Use it for an established TensorFlow estate or deployment path built around TF Serving and TFLite. |
| numba | PyPI | Use it to compile numerical Python loops without adopting JAX arrays, autodiff, or functional state. |
More ai / ml guides
openai · mcp · huggingface-hub · @modelcontextprotocol/sdk · scikit-learn · tiktoken · the whole shelf →
How this guide is made: grounded in the library's documentation, release notes, changelog, and issue history, on a fixed rubric — not a hands-on install of every release. The 50 most-downloaded entries are additionally install-verified in clean containers. Corrections: contact the desk.