JAX Cheat Sheet

1Install & jnp basics0.7.x

pip install "jax" # CPU pip install "jax[cuda12]" # NVIDIA GPU★
TPU on Cloud via jax[tpu]. Python 3.11+ required.
import jax, jax.numpy as jnp★
jnp mirrors the NumPy API. jax.numpy ops run on accelerators and are traceable by the transforms.
x = jnp.arange(12).reshape(3,4) · jnp.dot(a,b) · x.mean(axis=0)
Same calls as NumPy. Default float is float32; enable float64 with jax.config.update("jax_enable_x64", True).
jax.devices() · x.devices() · jax.device_put(x, dev)
List accelerators, see where an array lives, move it. Arrays are jax.Array regardless of backend.

2Immutable arrays & .at[]functional updates

x[0] = 1 # TypeError — arrays are immutablekey idea
You can't mutate a JAX array in place. Every "update" returns a new array.
y = x.at[0].set(1)★
The functional-update idiom. Also .add(v), .mul(v), .min(), .max(), .apply(fn).
y = x.at[idx].add(1) # scatter-add at integer indices
Out-of-bounds indices are dropped by default (configurable with mode=).
jnp.where(cond, a, b) · jnp.clip(x, lo, hi)
Vectorized selection instead of branching per element — the functional way to "modify conditionally".

3jit — compile with XLAspeed

@jax.jit def f(x, y): return jnp.sin(x) + y★
Traces the function once per input shape/dtype and compiles a fused XLA kernel. First call is slow; the rest are fast.
f = jax.jit(g, static_argnums=(1,))★
Mark args that control shapes/branches as static (e.g. a Python int used in reshape) — they trigger recompiles when changed.
jax.jit(f, donate_argnums=(0,))
Donate a buffer so XLA can reuse its memory in-place — important for large-model updates.
recompiles on new shapesgotcha
Every distinct input shape recompiles. Pad to fixed shapes for variable-length data; avoid Python-int args that vary each call.

4grad — automatic differentiationgradients

df = jax.grad(f) · df(3.0) # derivative wrt 1st arg★
f must return a scalar. argnums=(0,1) for multiple inputs; differentiate a pytree of params directly.
val, g = jax.value_and_grad(loss)(params, batch)★
Loss and its gradient in one pass — the standard training-step call.
jax.jacobian(f) · jax.hessian(f) · jax.jvp / jax.vjp
Full Jacobians/Hessians and the forward/reverse building blocks. has_aux=True to return extra state alongside the gradient.
@jax.jit def step(params, batch): g = jax.grad(loss)(params, batch) return jax.tree.map(lambda p,gr: p - 0.1*gr, params, g)★
jit + grad + pytree map = a complete SGD step, no framework needed.

5vmap — auto-vectorizebatch for free

batched = jax.vmap(f) · batched(X) # f written for one example★
Write the single-example function; vmap adds the batch dimension without a Python loop.
jax.vmap(f, in_axes=(0, None), out_axes=0)
Control which arg axes are mapped (None = broadcast/shared) and where the output axis lands.
jax.jit(jax.vmap(jax.grad(f)))★
Per-example gradients, batched, compiled — the composition that makes JAX special.
jax.lax.map(f, xs) # sequential map when vmap won't fit in memory
Falls back to a scan-based loop for memory-heavy cases.

6Randomness (explicit keys)no global seed

key = jax.random.key(0)★
The modern typed key (replaces PRNGKey). RNG state is a value you thread through — there is no hidden global seed.
key, sub = jax.random.split(key)★
Split before each draw; reuse a key and you get the same numbers. split(key, 4) for several subkeys.
jax.random.normal(sub, (3,4)) · uniform · randint · bernoulli★
Every sampler takes a key as its first argument — explicit, reproducible, and safe under jit/vmap.
never reuse a keygotcha
Two draws from the same key are perfectly correlated. Always split and pass a fresh subkey.

7Control flow (lax)under jit

jax.lax.cond(pred, true_fn, false_fn, operand)★
Data-dependent branch that survives tracing. A Python if on a traced value errors under jit.
carry, ys = jax.lax.scan(step_fn, init, xs)★
Compiled loop with accumulation (RNN/optimizer steps) — far faster to compile than an unrolled Python loop.
jax.lax.fori_loop(0, n, body, init) · while_loop(cond, body, init)
Counted / conditional loops with no Python-level unrolling.
python if/for on traced valuesgotcha
Branching on an array value under jit raises a TracerBoolConversionError — use lax.cond/jnp.where/lax.scan.

8Pytreesnested params

params = {"w": W, "b": b} # any nested dict/list/tuple of arrays★
A "pytree" is any nested container of arrays. Transforms operate on whole pytrees — model params are just a pytree.
jax.tree.map(lambda p, g: p - lr*g, params, grads)★
Apply a function leaf-wise across matching trees — the parameter-update idiom. (jax.tree_util for the fuller API.)
leaves, treedef = jax.tree.flatten(params) params = jax.tree.unflatten(treedef, leaves)
Flatten to a list and back — useful for serialization or custom optimizers.
jax.tree.reduce(operator.add, tree)
Aggregate across all leaves (e.g. total param count, global grad norm).

9Devices & shardingmulti-GPU/TPU

from jax.sharding import Mesh, NamedSharding, PartitionSpec as P★
The modern API (the jax.experimental.shard module was removed — use jax.sharding).
mesh = jax.make_mesh((4,2), ("data","model")) shard = NamedSharding(mesh, P("data", "model"))★
Describe device layout + how array axes map onto it. JAX 0.7 uses Shardy as the default partitioner.
x = jax.device_put(x, shard) · jax.jit(f)(x)★
Place data sharded, then a plain jit auto-parallelizes the computation — no explicit pmap needed.
jax.shard_map(f, mesh, in_specs=..., out_specs=...)
Explicit per-shard SPMD when you want manual control of collectives (the modern pmap successor).

10Flax / Optax / debugbuild real models

pip install flax # nn layers — flax.nnx (current API)★
JAX itself has no layers. Flax (nnx) provides modules/state; the older flax.linen API is still around.
pip install optax # optimizers opt = optax.adamw(1e-3); state = opt.init(params)★
Optax = composable gradient transforms. updates, state = opt.update(grads, state, params) then optax.apply_updates.
jax.debug.print("x={x}", x=x) · jax.debug.breakpoint()★
Printing/breakpoints that work inside jit (a bare print only runs at trace time).
with jax.disable_jit(): ... · jax.make_jaxpr(f)(x)
Run eagerly to debug, or inspect the traced program (jaxpr).

!The big fourmemorize these

side effects don't survive jit
A print/list-append/counter inside a jitted fn runs only at trace time, once. Use jax.debug.print and return values instead of mutating state.
no in-place mutation
x[i] = v fails; use x.at[i].set(v). Loops that "accumulate into" an array should use lax.scan.
python control flow on traced values
if x > 0 where x is traced → error. Use jnp.where, lax.cond, lax.scan.
reused / global RNG
No implicit seed — thread and split keys. Reusing a key gives identical draws.
recompilation blowupsperf
Varying shapes or non-static Python args recompile every call. Keep shapes fixed; mark shape-controlling args static_argnums.

JAX cheat sheet