PyPIAI / MLupdated 05 Aug 2026

jax

Item: jax
Rating: 4
Author: Keyoor

JAX is Google's library for accelerated array computing: you write NumPy-style Python, and JAX transforms it. jax.grad differentiates your function, jax.jit compiles it to fast machine code via the XLA compiler, and jax.vmap vectorizes it over a batch dimension. The transformations compose, so jit(vmap(grad(f))) is normal usage. It runs the same code on CPU, NVIDIA GPU, and Google TPU, and it is the foundation under training stacks like Flax and Optax.

jax · v0.11.0 · 4.9M/wk · ★ 36,108 · pushed 05 Aug 2026 · jax-ml/jax · docs · reviewed by Keyoor

Verdict

The best tool available for research-grade autodiff and TPU-scale training, and genuinely fun once the functional style clicks. Pick it for math-heavy work you write yourself; pick PyTorch when your job is mostly consuming the existing model ecosystem.

API stability3/5Still 0.x after eight years and it acts like it: deprecations land regularly (jax.random.PRNGKey to jax.random.key, tree_util reshuffles), and jaxlib version pinning is strict between releases.

Docs4/5docs.jax.dev has strong tutorials, an autodiff cookbook, and an honest Common Gotchas page, but advanced topics like custom_vjp and sharding assume you will read source and GitHub discussions.

Maintenance5/5Pushed the day this was written, with constant releases from the Google DeepMind team; the roughly 1,698 open issues reflect enormous surface area and traffic, not neglect.

Ecosystem4/5Flax, Optax, Orbax, and the scientific stack are healthy and TPU support is unmatched, but the pretrained-model and deployment ecosystem is far smaller than PyTorch's.

Use it if

You are doing ML research or scientific computing where you need gradients of arbitrary numerical code, including through loops and recursion
You want one codebase that runs on CPU, GPU, and TPU, especially TPU, where JAX is the best-supported path
Your workload is shaped like pure array math (physics sims, diffusion models, RL environments) and XLA fusion gives you real speedups over eager PyTorch or NumPy
You are scaling to many devices and want compiler-driven sharding (jax.jit with sharding annotations) instead of hand-written distributed code

Skip it if

You mainly fine-tune or serve existing models: the pretrained-model ecosystem (Hugging Face, vLLM, most tutorials) is PyTorch-first, and you will spend time porting weights instead of working
Your team is new to functional programming: JAX requires pure functions, immutable arrays (x.at[i].set(v) instead of x[i] = v), and explicit PRNG keys, and the tracer errors when you get this wrong are famously confusing
You need dynamic shapes or data-dependent control flow: jit retraces on every new shape and Python if statements on traced values fail, so ragged workloads fight the compiler constantly
You are on Windows with a GPU: NVIDIA GPU support on native Windows is not available, only WSL2 experimentally, and the wheel matrix (jax, jaxlib, CUDA plugins) is strict about matching versions
The README itself says it: this is a research project with sharp edges, not an official Google product, and 0.x versioning means APIs still move between releases

Setup reality

pip install -U jax gets you CPU in one line, but accelerators are where setup gets real: you install extras like jax[cuda13] or jax[tpu], and jaxlib plus the CUDA plugin packages must match the jax version exactly (0.11.0 pins jaxlib>=0.11.0,<=0.11.0). Python 3.12+ is required as of this release, which rules out older environments. Expect the first call to any jitted function to be slow because compilation happens on first trace, and expect to relearn habits: no in-place mutation, no side effects inside jit, and random numbers need explicit key management with jax.random.key and split.

Patterns

Differentiate a function with jax.gradcompute-gradient

import jax
import jax.numpy as jnp

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

grad_fn = jax.grad(loss)          # gradient w.r.t. first arg
g = grad_fn(w, x, y)
gw, gx = jax.grad(loss, argnums=(0, 1))(w, x, y)

grad differentiates with respect to argument 0 by default; the function must return a scalar or grad raises an error.

Compile a function with jax.jitjit-compile

import jax
import jax.numpy as jnp

@jax.jit
def step(x):
    return x * x + 2.0 * x

x = jnp.ones((5000, 5000))
step(x)                    # first call: traces and compiles (slow)
step(x).block_until_ready()  # later calls: fast

JAX dispatch is async; benchmark with block_until_ready() or you are timing only the dispatch, not the compute. New input shapes trigger recompilation.

Batch a per-example function with vmapvectorize-batch

import jax
import jax.numpy as jnp

def predict(w, x):        # x: single example, shape (d,)
    return jnp.dot(w, x)

batched = jax.vmap(predict, in_axes=(None, 0))
out = batched(w, xs)      # xs: (batch, d) -> out: (batch,)

in_axes=(None, 0) means broadcast w and map over axis 0 of xs. Composing vmap with grad gives per-example gradients without a Python loop.

Generate random numbers with explicit keysrandom-numbers

import jax

key = jax.random.key(42)
key, sub = jax.random.split(key)
x = jax.random.normal(sub, (3, 3))
keys = jax.random.split(key, 8)   # one key per parallel draw

Reusing a key gives identical numbers, the top gotcha for newcomers. Split before every draw. jax.random.key replaced the older PRNGKey API.

Get loss and gradient in one passvalue-and-grad

loss_and_grad = jax.value_and_grad(loss)
loss_val, grads = loss_and_grad(params, x, y)

# with auxiliary outputs:
(loss_val, metrics), grads = jax.value_and_grad(
    loss_with_metrics, has_aux=True)(params, x, y)

Cheaper than calling the function and grad separately because the forward pass is shared. has_aux lets you return metrics without breaking the scalar-output rule.

Apply an SGD update across a parameter treeupdate-pytree-params

import jax

grads = jax.grad(loss)(params, x, y)
params = jax.tree.map(
    lambda p, g: p - 0.01 * g, params, grads)

Any nest of dicts, lists, and tuples of arrays is a pytree; jax.tree.map walks matching structures. jax.tree.map is the current spelling of the old jax.tree_util.tree_map.

Update array elements without mutationfunctional-array-update

import jax.numpy as jnp

x = jnp.zeros((5,))
y = x.at[2].set(7.0)      # x is unchanged
z = x.at[1:3].add(1.0)
w = x.at[0].max(3.0)

JAX arrays are immutable, so x[2] = 7 raises an error. Under jit the .at updates compile to in-place operations, so there is no copy cost.

Mark configuration arguments as staticjit-static-args

from functools import partial
import jax

@partial(jax.jit, static_argnums=(1,))
def rollout(state, num_steps):
    for _ in range(num_steps):   # Python loop, unrolled at trace time
        state = step(state)
    return state

Values used in Python control flow or shapes must be static. Every new static value compiles a new specialization, so do not pass unbounded variety.

Write a compiled loop with lax.scanscan-loop

import jax
from jax import lax

def step(carry, x):
    carry = carry * 0.9 + x
    return carry, carry        # (new_carry, per-step output)

final, history = lax.scan(step, init=0.0, xs=data)

scan compiles the loop body once instead of unrolling, keeping compile times sane for long sequences. It is the idiomatic way to write RNNs and optimizer loops.

Print traced values inside a jitted functiondebug-inside-jit

import jax

@jax.jit
def f(x):
    y = x + 1
    jax.debug.print("y = {y}", y=y)
    return y * 2

# or disable jit entirely while debugging:
with jax.disable_jit():
    f(3.0)

A normal print inside jit fires once at trace time and shows tracers, not values. jax.debug.print prints at runtime; disable_jit gives you plain eager Python.

Place data on devices with an explicit meshshard-across-devices

import jax
from jax.sharding import Mesh, 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)  # compiler parallelizes

jit propagates shardings automatically, so annotating inputs is often all a data-parallel job needs. Print jax.typeof(x) to see how an array is sharded.

Alternatives

Package	Registry	Pick it when
torch	PyPI	You want the dominant ecosystem, eager-mode debugging, and pretrained models that load without conversion.
tensorflow	PyPI	You are tied to TF serving, TFLite, or an existing TF codebase; for new work most teams pick JAX or PyTorch instead.
numba	PyPI	You just want to JIT-compile numeric Python loops without adopting a new array library or autodiff.