JAX for PyTorch Users

How it works

When you call a jit-wrapped function for the first time with a given combination of input shapes/dtypes, JAX:

1. Runs your Python function once with tracers (abstract stand-ins that only carry shape/dtype, not concrete values) instead of real arrays.
2. Records every JAX operation performed into an intermediate representation called a jaxpr.
3. Hands the jaxpr to XLA, which compiles it into optimized machine code for your target device (CPU/GPU/TPU).
4. Caches the compiled artifact, keyed by the function identity and the input shapes/dtypes.

Subsequent calls with the same shapes/dtypes skip tracing and compilation entirely and just execute the cached compiled program — this is where the speedup comes from.

import jax
import jax.numpy as jnp

def selu(x, alpha=1.67, lambda_=1.05):
    return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = jnp.arange(1_000_000.0)

selu_jit = jax.jit(selu)
selu_jit(x).block_until_ready()   # first call: traces + compiles (slow)
selu_jit(x).block_until_ready()   # subsequent calls: just runs (fast, ~8x here)

block_until_ready() exists because JAX dispatches work asynchronously — a call returns a future-like array immediately while the device keeps computing, similar in spirit to CUDA kernel launches being asynchronous in PyTorch, except JAX makes this explicit for benchmarking rather than hiding it behind a torch.cuda.synchronize() you have to remember.

The decorator form is more common in practice:

@jax.jit
def selu(x, alpha=1.67, lambda_=1.05):
    return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

Why Python control flow breaks under jit

Because tracing replaces your inputs with abstract tracers that only know shape/dtype (not concrete values), any Python control flow that needs a concrete value — if x > 0:, while i < n: — cannot be resolved during tracing:

def f(x):
    if x > 0:          # x is a Tracer here, not a real number
        return x
    else:
        return 2 * x

jax.jit(f)(10)   # TracerBoolConversionError

def g(x, n):
    i = 0
    while i < n:        # n is also a Tracer
        i += 1
    return x + i

jax.jit(g)(10, 20)   # TracerBoolConversionError

There are two ways to fix this, and picking the right one matters:

Option A — mark the offending argument as static. static_argnums/static_argnames tell jit to treat that argument as a literal Python value baked into the compiled program (a new compilation happens for each distinct value), rather than a traced array:

f_jit = jax.jit(f, static_argnums=0)
f_jit(10)   # works: 10

g_jit = jax.jit(g, static_argnames=['n'])
g_jit(10, 20)   # works: 30

Use this sparingly — every distinct static value triggers a fresh recompilation, so this only scales to a small, bounded set of values (e.g. a training: bool flag, a small integer config), never something like a batch loss value.

Option B — use a structured control-flow primitive (lax.cond, lax.while_loop, lax.fori_loop, lax.scan) so the branch/loop itself becomes part of the compiled program instead of something Python needs to resolve eagerly. See Section 11.

A caching gotcha

jit caches by function identity. Defining a new closure inside a loop defeats the cache and forces a recompile every iteration — a subtle performance bug that has no PyTorch analogue, since PyTorch doesn’t cache compiled graphs by function object identity in the same way:

# BAD: a fresh function object every iteration -> recompiles every time
def g_bad(x, n):
    i = 0
    while i < n:
        i = jax.jit(lambda i: i + 1)(i)   # new lambda -> new cache entry -> recompile
    return x + i

# GOOD: reuse one compiled function
_step = jax.jit(lambda i: i + 1)
def g_good(x, n):
    i = 0
    while i < n:
        i = _step(i)
    return x + i

Compare to PyTorch: torch.compile(model) gives you similar shape-specialized caching under the hood, but eager PyTorch (no compile) never has this failure mode at all, because every op just runs immediately with concrete tensors. The tradeoff is that eager PyTorch pays Python dispatch overhead on every single op, every single call.

Higher-order derivatives

Because jax.grad returns an ordinary function, and ordinary functions can be differentiated again, higher-order derivatives are just repeated application:

f = lambda x: x**3 + 2*x**2 - 3*x + 1

dfdx   = jax.grad(f)
d2fdx  = jax.grad(dfdx)
d3fdx  = jax.grad(d2fdx)

print(dfdx(1.0))    # 4.0
print(d2fdx(1.0))   # 10.0
print(d3fdx(1.0))   # 6.0

In PyTorch this requires torch.autograd.grad(..., create_graph=True) chained manually — a noticeably more verbose pattern for the same idea.

Gradients with respect to specific / multiple arguments

Use argnums to choose which positional argument(s) to differentiate with respect to — this replaces the PyTorch pattern of selectively setting requires_grad=True on some tensors and not others:

def loss(W, b, x, y):
    preds = jnp.dot(x, W) + b
    return jnp.mean((preds - y) ** 2)

W_grad = jax.grad(loss)(W, b, x, y)              # grad wrt W only (argnums=0 is default)
W_grad, b_grad = jax.grad(loss, argnums=(0, 1))(W, b, x, y)   # grad wrt both

value_and_grad: getting the loss and the gradient together

Training loops always need both the loss value (to log) and the gradient (to step the optimizer). Calling grad alone throws away the forward value; jax.value_and_grad computes both in a single pass — this is the direct equivalent of PyTorch’s loss = criterion(...); loss.backward() where you already have loss sitting around for free:

loss_value, (W_grad, b_grad) = jax.value_and_grad(loss, argnums=(0, 1))(W, b, x, y)

Gradients of nested structures (dicts, NamedTuples, …)

jax.grad works transparently over pytrees (see Section 8) — you don’t need every parameter to be a flat positional argument:

params = {'W': W, 'b': b}

def loss2(params, x, y):
    preds = jnp.dot(x, params['W']) + params['b']
    return jnp.mean((preds - y) ** 2)

grads = jax.grad(loss2)(params, x, y)
# grads == {'W': ..., 'b': ...}, same structure as params

This is exactly how gradients of full neural-net parameter pytrees work in Section 13 — no .parameters() iterator needed, because the “parameters” are just a regular (if deeply nested) Python data structure.

Jacobians, Hessians, and stop_gradient

For vector-valued functions, jax.grad alone isn’t enough (it requires a scalar output). Use jax.jacrev / jax.jacfwd for full Jacobians, and compose them for Hessians:

def f(x):
    return jnp.array([jnp.sin(x[0]) * x[1], x[0] * x[1] ** 2])

jac = jax.jacrev(f)(jnp.array([1.0, 2.0]))   # reverse-mode: efficient for few inputs, many outputs
jac = jax.jacfwd(f)(jnp.array([1.0, 2.0]))   # forward-mode: efficient for many inputs, few outputs

hessian = jax.jacfwd(jax.jacrev(f))   # standard JAX idiom for a Hessian

To exclude part of a computation from the gradient (PyTorch’s .detach()), use jax.lax.stop_gradient:

def loss_fn(params, x, target):
    pred = model(params, x)
    target = jax.lax.stop_gradient(target)   # equivalent to target.detach() in PyTorch
    return jnp.mean((pred - target) ** 2)

Per-example gradients: where grad + vmap shine

A pattern that’s awkward in PyTorch (historically requiring functorch/torch.func.vmap) is trivial by composing two JAX transformations:

# grad of a single example's loss, then vmap it over a batch
per_example_grads = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))(params, x_batch, y_batch)

params is broadcast (in_axes=None), while x_batch/y_batch are mapped over their leading axis — giving you one gradient per example instead of one gradient averaged over the batch. This is the same vmap covered next, just composed with grad.

in_axes and out_axes

in_axes tells vmap which axis of each argument is the batch axis (None means “this argument is not batched, broadcast it as-is”):

# batch dimension is axis 1 instead of axis 0, for both inputs and the output
auto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)
auto_batch_convolve_v2(xst, wst)

# only `xs` is batched; the same `w` is reused for every example
batch_convolve_v3 = jax.vmap(convolve, in_axes=[0, None])
batch_convolve_v3(xs, w)

Compare to PyTorch: this replaces two different habits at once — (1) manually rewriting a function to insert a batch dimension everywhere, and (2) torch.func.vmap (formerly functorch.vmap), which is a bolt-on, less central feature in PyTorch. In JAX, vmap is a first-class, everyday transformation, and it composes with jit and grad freely: jax.jit(jax.vmap(jax.grad(f))) is completely ordinary JAX code.

Splitting keys

Since reusing a key gives you the same “random” numbers again, and there’s no global state to “advance,” you explicitly split a key into fresh, independent subkeys whenever you need new randomness — treat every key as single-use:

key = random.key(0)

key, subkey = random.split(key)
val1 = random.normal(subkey)     # use subkey once, then discard it

key, subkey = random.split(key)
val2 = random.normal(subkey)     # `key` keeps getting split for each new draw

# need several keys at once (e.g. one per layer during init)?
key, *subkeys = random.split(key, num=5)

A useful mental model: key is like a seed you’re not allowed to reuse. split is how you turn one seed into two (or more) independent seeds — one to consume right now, and one (conventionally the first one returned) to keep threading through the rest of your program.

Non-obvious detail: individually sampling from split subkeys does not generally match a single vectorized call with a shape argument — don’t expect [random.normal(k) for k in random.split(key, 3)] to equal random.normal(key, shape=(3,)). Both are valid, reproducible random streams; they’re just different streams.

Compare to PyTorch: this replaces torch.manual_seed(0) plus relying on the global generator advancing correctly across your whole program (which quietly breaks under DataLoader multiprocessing unless you’re careful with worker_init_fn, and is not something you can trivially reproduce on TPU pods). In JAX, correct/reproducible parallel randomness is what you get by construction, at the cost of manually threading key/subkey everywhere you’d have called a torch.rand* function.

10.1 In-place updates don’t exist

Already covered in Section 3 — x[i] = v raises an error; use x.at[i].set(v) (and .add(), .multiply(), .min(), .max(), .get()).

10.2 Out-of-bounds indexing is clamped, not an error

jnp.arange(10)[11]   # returns 9 (index clamped to the last valid position) — no IndexError!

This is deliberate: on accelerators, raising exceptions from out-of-bounds indices inside a compiled kernel isn’t well defined, so JAX defines it as clamping instead of erroring. If you need NaN-on-out-of-bounds behavior for debugging, use .at[].get(mode='fill', fill_value=jnp.nan):

jnp.arange(10.0).at[11].get(mode='fill', fill_value=jnp.nan)   # nan

PyTorch equivalent: torch.tensor([...])[11] raises IndexError. Coming from PyTorch, silently-wrong-instead-of-crashing is the dangerous direction to be surprised in — add your own bounds assertions during debugging if this matters for your code.

10.3 Lists and tuples of numbers aren’t arrays

jnp.sum([1, 2, 3])         # TypeError in strict contexts / silently slow in others — don't do this
jnp.sum(jnp.array([1, 2, 3]))   # correct

NumPy silently converts Python lists to arrays for you; JAX would rather you be explicit, since implicit conversion inside a traced function can quietly turn into a very slow, unrolled Python-level operation.

10.4 Dynamic shapes are not allowed under jit

Every array’s shape must be static (known at trace time) — a shape cannot depend on the value of an array, only on its shape/dtype:

def nansum(x):
    mask = ~jnp.isnan(x)
    x_without_nans = x[mask]   # shape of the result depends on how many NaNs are in x -> fails under jit
    return x_without_nans.sum()

The fix is a very common JAX idiom: mask with jnp.where instead of shrinking the array, so the shape stays fixed and only the values change:

@jax.jit
def nansum_fixed(x):
    mask = ~jnp.isnan(x)
    return jnp.where(mask, x, 0).sum()   # same shape in and out, NaNs replaced by 0 before summing

PyTorch comparison: boolean masking (x[mask]) works fine in eager PyTorch because nothing needs a static shape ahead of time. Under torch.compile, you’ll hit a close analogue of this same restriction (torch._dynamo guards / graph breaks on data-dependent shapes).

10.5 float32 by default, even for “double precision” literals

x = jnp.array([1.0, 2.0], dtype=jnp.float64)
x.dtype   # dtype('float32')  <-- silently downcast!

JAX defaults to 32-bit floats everywhere (matching typical accelerator throughput/memory tradeoffs) and silently truncates 64-bit inputs unless you opt in explicitly, at process startup, before any JAX arrays are created:

import jax
jax.config.update("jax_enable_x64", True)
# or set the environment variable JAX_ENABLE_X64=1 before importing jax

PyTorch comparison: torch.get_default_dtype() is also float32 by default, but PyTorch does not silently downcast an explicitly-constructed float64 tensor — torch.tensor([1.0], dtype=torch.float64).dtype stays float64. This is one of the sharper surprises for newcomers; if a metric or loss looks suspiciously imprecise, check whether it silently landed in float32.

10.6 Type promotion and casting rules differ from NumPy

np.arange(254.0, 258.0).astype('uint8')    # [254, 255, 0, 1]      (numpy wraps around)
jnp.arange(254.0, 258.0).astype('uint8')   # [254, 255, 255, 255]  (jax clips, and this is backend-dependent)

Unsafe casts near dtype boundaries are not guaranteed to match NumPy exactly, and can vary between CPU/GPU/TPU. Don’t rely on overflow-wraparound behavior being portable.

10.7 JIT + class methods

Wrapping a bound method (self.foo) with jax.jit naively re-traces every time self changes identity (e.g. a new instance each call), and — per the pure-functions rule — mutating self inside the method silently doesn’t “stick” the way you’d expect from PyTorch’s nn.Module. Three common fixes: (1) pull the logic into a free function and jit that instead, (2) mark self as static via static_argnums if it’s hashable and small, or (3) register the class itself as a pytree. In practice, Flax NNX (Section 13) handles this for you, which is the main reason to reach for it instead of hand-rolling classes.

10.8 Debugging NaNs/Infs

jax.config.update("jax_debug_nans", True)   # raise as soon as a NaN is produced, with a traceback to the op
jax.config.update("jax_debug_infs", True)

This is the closest JAX equivalent to torch.autograd.detect_anomaly() — turn it on only while debugging, since it disables some optimizations and slows execution.

Worked example: linear regression parameters as state

This is precisely the pattern a training loop uses, with model parameters standing in for self.n:

from typing import NamedTuple
import jax
import jax.numpy as jnp

class Params(NamedTuple):
    weight: jnp.ndarray
    bias: jnp.ndarray

def loss(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
    pred = params.weight * x + params.bias
    return jnp.mean((pred - y) ** 2)

LEARNING_RATE = 0.005

@jax.jit
def update(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> Params:
    grad = jax.grad(loss)(params, x, y)
    # jax.tree.map applies the SGD update to every leaf (weight AND bias) at once
    new_params = jax.tree.map(lambda p, g: p - g * LEARNING_RATE, params, grad)
    return new_params

params = Params(weight=jnp.zeros(()), bias=jnp.zeros(()))
for _ in range(1000):
    params = update(params, xs, ys)   # params is reassigned, never mutated in place

# The PyTorch equivalent — mutation is implicit and hidden inside .step()
model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)

for _ in range(1000):
    optimizer.zero_grad()
    pred = model(xs)
    loss = ((pred - ys) ** 2).mean()
    loss.backward()
    optimizer.step()   # mutates model.parameters() in place, invisibly

Neither version is “more correct” — but notice that the JAX version makes every state transition visible in the return value and the reassignment (params = update(...)), while PyTorch’s happens as a side effect buried inside .step(). This explicitness is exactly what lets jit, grad, and vmap reason about — and safely transform — the whole computation. Flax’s nnx.Optimizer (next section) gives you back something that feels like optimizer.step(), while still being built on this explicit-state foundation underneath.

Why NNX and not “plain JAX classes”?

Flax NNX modules hold their parameters as ordinary Python attributes (self.w = nnx.Param(...)) — which looks exactly like the PyTorch self.weight = nn.Parameter(...) habit — while internally staying compatible with JAX’s pytree/functional model, so nnx.jit and nnx.grad can still trace through them safely. It removes the two sharp edges from Section 10.7 and Section 12 without requiring you to manually thread every parameter through function arguments by hand.

Defining a layer

from flax import nnx
import jax.numpy as jnp
import jax

class Linear(nnx.Module):
    def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
        self.w = nnx.Param(rngs.params.uniform((din, dout)))
        self.b = nnx.Param(jnp.zeros((dout,)))
        self.din, self.dout = din, dout

    def __call__(self, x: jax.Array):
        return x @ self.w + self.b[None]

model = Linear(din=2, dout=5, rngs=nnx.Rngs(params=0))
y = model(jnp.ones((1, 2)))

# The PyTorch equivalent
import torch.nn as nn

class Linear(nn.Module):
    def __init__(self, din, dout):
        super().__init__()
        self.w = nn.Parameter(torch.rand(din, dout))
        self.b = nn.Parameter(torch.zeros(dout))

    def forward(self, x):
        return x @ self.w + self.b

model = Linear(2, 5)
y = model(torch.ones(1, 2))

The differences to flag for a PyTorch reader: NNX requires you to pass an explicit rngs: nnx.Rngs for parameter initialization (there’s no hidden global RNG — see Section 9), and shapes (din, dout) must be given up front rather than lazily inferred on first call (there’s no nn.LazyLinear equivalent).

Composing modules (an MLP)

class MLP(nnx.Module):
    def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
        self.linear1 = Linear(din, dmid, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.1, deterministic=False)
        self.bn = nnx.BatchNorm(dmid, use_running_average=False, rngs=rngs)
        self.linear2 = Linear(dmid, dout, rngs=rngs)

    def __call__(self, x: jax.Array, rngs: nnx.Rngs):
        x = self.linear1(x)
        x = self.bn(x)
        x = self.dropout(x, rngs=rngs)
        x = nnx.gelu(x)
        return self.linear2(x)

model = MLP(din=2, dmid=16, dout=5, rngs=nnx.Rngs(0))
y = model(jnp.ones((3, 2)), rngs=nnx.Rngs(1))

Submodules are just attributes, composed exactly the way nn.Sequential/nested nn.Modules are in PyTorch — this part of the mental model transfers almost directly.

Optimizers with Optax

Optax optimizers are stateless transformations of gradients, not stateful objects that own your parameters. The core pattern (used with or without Flax) is:

import optax

optimizer = optax.adam(learning_rate=1e-3)
opt_state = optimizer.init(params)              # params: any pytree of arrays

grads = jax.grad(loss_fn)(params, x, y)
updates, opt_state = optimizer.update(grads, opt_state)   # opt_state is explicit, threaded state
params = optax.apply_updates(params, updates)             # apply the updates functionally

# The PyTorch equivalent
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

optimizer.zero_grad()
loss = loss_fn(model, x, y)
loss.backward()
optimizer.step()    # optimizer owns and mutates model.parameters() in place

optax.adam(...) replaces torch.optim.Adam(...), and Optax has direct equivalents of nearly everything in torch.optim: optax.sgd, optax.adamw, optax.rmsprop, learning-rate schedules (optax.cosine_decay_schedule, optax.warmup_cosine_decay_schedule), and gradient transformations you compose with optax.chain (e.g. gradient clipping via optax.clip_by_global_norm) — the Optax equivalent of stacking multiple torch.optim tricks by hand.

nnx.Optimizer: wiring model + optimizer together

Flax NNX wraps this pattern so it feels like optimizer.step(), while still being pure functional Optax underneath:

model = MLP(din=2, dmid=16, dout=10, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

@nnx.jit
def train_step(model, optimizer, x, y, rngs):
    def loss_fn(model: MLP, rngs: nnx.Rngs):
        y_pred = model(x, rngs)
        return jnp.mean((y_pred - y) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model, rngs)
    optimizer.update(model, grads)   # updates model's params in place (NNX allows this internally)
    return loss

x, y = jnp.ones((5, 2)), jnp.ones((5, 10))
loss = train_step(model, optimizer, x, y, nnx.Rngs(2))

nnx.jit and nnx.grad/nnx.value_and_grad are NNX-aware wrappers around jax.jit/jax.grad — they know how to split an NNX module into its underlying pytree of parameters before tracing and reassemble it afterward, which is what makes optimizer.update(model, grads) safe to call and have it “stick,” unlike the raw self.n += 1 example from Section 12. Under the hood it is still following the pure state-in/state-out contract; NNX is just doing the threading for you.

Meshes and PartitionSpec

A mesh names your physical devices along one or more logical axes; a PartitionSpec (jax.P) then describes, per array axis, which mesh axis (if any) it’s split across:

import jax
import jax.numpy as jnp

jax.set_mesh(jax.make_mesh((4, 2), ('X', 'Y')))   # 8 devices arranged as a 4x2 logical grid

x = jnp.arange(32.0).reshape(8, 4)
x = jax.device_put(x, jax.P('X', 'Y'))    # axis 0 split across mesh axis 'X', axis 1 across 'Y'

print(jax.typeof(x))    # float32[8@X,4@Y]

y = jax.device_put(x, jax.P('Y', 'X'))          # a different split
z = jax.device_put(x, jax.P('X', None))         # split over 'X', replicated over 'Y'

jax.jit auto-parallelizes based on sharding

Once inputs carry a sharding, ordinary jit-compiled code is automatically partitioned across devices — reductions like .sum() transparently become all-reduces:

@jax.jit
def compute(x):
    y = x.sum(0)     # the compiler inserts an all-reduce across the sharded axis automatically
    return jnp.sin(y)

x = jax.device_put(jnp.arange(8.0), jax.P('X'))
result = compute(x)

Data parallelism vs. model parallelism, in one line each

# Data parallelism: shard the batch dimension across devices (like DistributedDataParallel)
x = jax.device_put(jnp.ones((batch_size, seq_len)), jax.P('X', None))

# Model/tensor parallelism: shard a weight matrix across devices (like tensor-parallel FSDP-style sharding)
params = jax.device_put(jnp.ones((model_dim, model_dim)), jax.P('Y', None))

Because both are just Shardings on ordinary jax.Arrays, you can mix data and model parallelism on the same mesh (e.g. ('X', 'Y') = (data axis, model axis)) without separate parallelism frameworks for each.

Manual control: shard_map

When you want to write per-device code explicitly (and issue explicit collectives yourself, rather than letting the compiler infer them), jax.shard_map gives you that — each device runs your function on just its local shard:

mesh = jax.make_mesh((4, 2), ('X', 'Y'))
jax.set_mesh(mesh)

@jax.shard_map(out_specs=jax.P('X', None))
def manual_matmul(x_shard, y_shard):
    z_partial = jnp.dot(x_shard, y_shard)
    return jax.lax.psum_scatter(z_partial, 'X', tiled=True)   # explicit collective

result = manual_matmul(x, y)

Compare to PyTorch: this is roughly the JAX analogue of choosing between torch.compile + DTensor (automatic, sharding-annotation-driven parallelism, closest to jit + Sharding) versus hand-writing torch.distributed collectives yourself inside a custom parallel module (closest to shard_map). If you only need standard data parallelism across a handful of GPUs, jax.device_put(..., jax.P('X', ...)) plus ordinary jit gets you most of the way with far less code than either PyTorch approach.