JAX for PyTorch Researchers

You open a paper's reference repository to change one loss term. The code looks like NumPy, but parameters live in nested dictionaries, random-number keys move through every function, and decorators wrap decorators. Rewriting the project in PyTorch feels tempting. A smaller move is usually safer: learn the execution contract, change the local function, and keep the paper's tests and checkpoints intact.

You already know tensors, gradients, and training loops from PyTorch. JAX keeps that mathematics, but moves ownership into function inputs and outputs. Once state and shapes become visible, unfamiliar code starts reading like a precise dataflow program instead of implicit framework behavior.

Read the program before choosing a rewrite

JAX is an array-computing and program-transformation library. jax.numpy supplies a NumPy-shaped interface, while transformations such as grad, vmap, and jit derive new functions from functions you wrote.^{[1]Reference 1Quickstart: How to think in JAXhttps://docs.jax.dev/en/latest/quickstart.html} The central unit isn't a model object. It's a numerical function with an explicit input and output contract.

PyTorch researchers often meet JAX through a larger stack such as Flax, Equinox, Optax, or a lab-specific framework. Those libraries may package state differently, but core JAX still determines how arrays are traced, differentiated, vectorized, compiled, and dispatched. Read through library wrappers until you find the pure numerical boundary.

The first review pass should answer four questions:

1. Which values change after one step?
2. Which arguments control shapes or Python branches?
3. Where does each random key come from, split, and leave?
4. Which functions are transformed by grad, vmap, or jit?

That inventory often reveals a narrow edit: place a new regularizer in loss function, change sampling where a one-use key enters, or derive a per-example statistic with vmap instead of porting frameworks.

The comparison below maps familiar habits without pretending that either framework has only one programming style.

One step is a state transition

A PyTorch training loop often expresses state through object mutation. Model parameters change when the optimizer steps. Gradient buffers appear on tensors. A random generator advances behind dropout or sampling calls. You can inspect each object, but its changing state isn't written in the step's return type.

JAX transformations work best when changing values become ordinary data. A training step accepts current parameters, optimizer state, batch, and random key. It returns updated parameters, updated optimizer state, a retained next key, and metrics. Official JAX guidance uses this state-in, state-out pattern for model parameters, optimizer state, and stateful layer statistics.^{[2]Reference 2Stateful computationshttps://docs.jax.dev/en/latest/stateful-computations.html}

The figure makes ownership visible across one call.

A useful signature looks like this:

That snippet is structural pseudocode, not a promise about one optimizer library. Its value is the boundary. Nothing inside needs to reach into a global model, global key, or hidden optimizer object.

Pure means observable through arguments and returns

A pure function produces outputs from inputs without reading or writing hidden external state. JAX's transformation guides require this discipline because tracing may run Python code once, cache its numerical program, and skip Python side effects on later calls.^{[3]Reference 3JAX: The Sharp Bitshttps://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html}

Printing inside a jitted function illustrates the mismatch. Python print runs while JAX traces, so it may display a tracer once instead of a device value on every execution. Use returned metrics, jax.debug.print for deliberate transformed-code debugging, or inspect values outside the compiled boundary.

Globals are more dangerous. If a loss reads weight_decay from a module-level variable, a trace can capture its old value. Updating the Python global doesn't guarantee the cached executable changes. Pass research settings as arguments, then decide whether each one is an array value or a small static configuration.

Immutable arrays remove aliasing surprises

NumPy and PyTorch let code assign into an array or tensor. JAX's ordinary arrays use immutable value semantics. x[i] = value raises an error, and x.at[i].set(value) describes an updated array value instead.^{[3]Reference 3JAX: The Sharp Bitshttps://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html}

The old value remains valid from Python's perspective. The caller must bind the returned value. Inside compiled code, the compiler may reuse buffers when safe, but that optimization doesn't change the program's value semantics.

This runnable check updates one element and proves that the original array didn't change.

An easy porting mistake drops the return value:

The corrected version keeps the new leaf and then keeps the new tree:

For updates across every parameter leaf, jax.tree.map is clearer than rebuilding dictionaries by hand. Later update examples use it after introducing pytrees.

Randomness is explicit data

PyTorch code can seed a generator and let calls advance its internal state. JAX's standard random API passes a pseudo-random number generator (PRNG) key into each random function. Sampling doesn't mutate the key. Calling a random function twice with the same key produces the same sample and may create unwanted correlations.^{[4]Reference 4Pseudorandom numbershttps://docs.jax.dev/en/latest/random-numbers.html}

Split once for each independent use. Keep one output as future state and consume the others as one-use subkeys. A loop, batch, device, or model layer can receive its own derived key without synchronizing a hidden global generator.

This runnable example verifies both properties: key reuse repeats a sample, while separate subkeys produce different samples.

Typed keys from jax.random.key() are the current standard interface. You may still see older repositories call jax.random.PRNGKey(). Read the repository's key format before changing checkpoint or serialization code.

Split at ownership boundaries

Suppose one training step needs dropout and data augmentation. Split the step key into two subkeys, then pass each to one operation. Don't let both helpers split the same parent independently, since they'll derive matching child sequences.

Across devices, derive keys from stable identifiers such as process, device, step, or example indices when reproducibility requires that mapping. fold_in can combine an integer identifier with a key without managing a long manual split chain. Write the ownership rule beside the code so later vectorization doesn't duplicate random streams accidentally.

Current JAX also has experimental stateful RNG support. Its documentation still recommends explicit key semantics for performance-sensitive applications and names transformation limitations. This lesson keeps explicit keys because their dataflow remains visible under jit, vmap, checkpointing, and review.

Transform one small function at a time

JAX's transformations become easier when the base function is boring. Start with a scalar loss over one parameter tree and one example. Check its output. Differentiate it. Add a batch axis with vmap. Compile the stable outer call with jit only after eager checks pass.

The flow below shows function derivation, not runtime stages. Each box still represents a callable function.

jax.grad(f) returns a new function that evaluates the gradient of scalar-valued f. jax.value_and_grad(f) returns both scalar value and gradient, which avoids expressing the forward loss twice in a training step.^{[5]Reference 5Automatic differentiationhttps://docs.jax.dev/en/latest/automatic-differentiation.html}

jax.vmap(f) adds mapped array axes to a function. in_axes tells JAX which arguments carry the batch dimension and which stay shared. For a per-example loss, in_axes=(None, 0, 0) means one shared parameter tree plus batched inputs and targets.^{[6]Reference 6Automatic vectorizationhttps://docs.jax.dev/en/latest/automatic-vectorization.html}

jax.jit(f) stages compatible array work for compilation and caches a specialized executable. It doesn't make arbitrary Python dynamic. Lists that grow, file reads, network calls, and Python branches on traced values need to stay outside or be expressed through JAX-supported operations.^{[7]Reference 7Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

Transformation order changes the derived program

These expressions are both legal but ask for different derived functions:

The first returns one gradient tree per example. The second returns the gradient of a scalar batch mean. Averaging the first set of gradients can match the second for a simple mean of independent examples, but memory layout and intermediate materialization may differ. Match the transformation to the artifact you need.

Per-example gradients help with clipping, influence analysis, and variance inspection. Ordinary training usually needs the gradient of an aggregated scalar loss. Don't materialize a batch of full parameter-sized gradients merely because vmap makes the expression short.

Pytrees carry structured state

Real model state isn't one matrix. It can be a nested dictionary of layer weights, tuples of optimizer moments, and auxiliary statistics. JAX calls a nested container of leaves a pytree. Built-in containers such as dictionaries, tuples, and lists can form tree nodes, while arrays usually form leaves.^{[8]Reference 8Pytreeshttps://docs.jax.dev/en/latest/pytrees.html}

Transformations flatten the structure at their boundary, operate on leaves, and rebuild matching output structure. If params is a pytree, jax.grad(loss_fn)(params, batch) returns gradients with the same tree structure. That structural agreement is a powerful debugging invariant.

Use jax.tree.map to apply the same update over matching leaves:

Two trees passed to tree.map need compatible structure. A missing dictionary key, tuple with different length, or leaf shape mismatch should be treated as a state-schema bug, not patched with position-dependent indexing.

Custom model classes require care. Unless a class is registered or supported by its library as a pytree node, JAX may treat the entire object as one leaf. Before editing serialization or transformation boundaries, inspect jax.tree.structure(params) and jax.tree.map(lambda x: x.shape, params).

Separate dynamic leaves from static metadata

Array leaves can change value without changing the tree definition. Layer names, activation choices, and architecture sizes are often static metadata. Frameworks may place them in dataclass fields, auxiliary pytree data, closures, or static arguments.

Changing tree structure or static metadata can trigger a new trace or compilation. Changing only values inside same-shaped array leaves usually reuses the existing compiled variant. Keep architecture decisions outside hot per-step data.

Checkpoint conversion must preserve tree paths, leaf shapes, dtypes, and semantic layout. A matching leaf count isn't enough. A transposed projection matrix has same element count and still changes model behavior.

Translate a tiny training step

The familiar PyTorch loop makes mutation explicit through method calls. The snippet below performs one linear-regression step and shows where gradients and optimizer state live.

The JAX translation starts by exposing model parameters as arrays. Prediction and loss are plain functions. The step asks value_and_grad for a gradient pytree, maps an update across leaves, and returns the new parameter tree.

The complete CPU-sized example trains a two-feature linear model. It also uses vmap to inspect final per-example losses and blocks on the reported value before printing.

The update has no optimizer state because plain stochastic gradient descent only needs a learning rate here. Adam would add first moments, second moments, and a step counter to the inputs and outputs. An optimizer library packages that state, but the state transition remains.

Randomized training adds a key. Split it in train_step, give the one-use subkey to dropout or augmentation, and return the retained key beside parameters. Don't close over a module-level key.

JIT compiles a shape-specific contract

Just-in-time (JIT) compilation has an up-front cost. On the first compatible call, JAX traces the Python function, lowers array operations, compiles an executable, and runs it. Later calls can reuse that executable when the relevant input contract matches.^{[7]Reference 7Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

Input shapes and dtypes commonly participate in specialization. A [32, 512] token batch and a [16, 512] final batch can produce different compiled variants. So can float32 versus bfloat16 leaves. Static arguments add their concrete values to the cache key and recompile when those values change.

Control shape churn before optimizing kernels:

• pad batches to a fixed shape and carry a mask;
• bucket sequence lengths into a small set;
• drop or separately handle short remainder batches;
• keep static option values bounded;
• construct the jitted callable once outside hot loops;
• log or profile compilation separately from execution.

Tracers aren't ordinary Python values

During tracing, an array argument may become a tracer that records abstract properties such as shape and dtype. Python can't always turn that tracer into a concrete bool, int, or array index. A branch such as if loss < threshold: can raise TracerBoolConversionError inside jit.^{[7]Reference 7Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

Choose a fix by semantics:

Marking every awkward value static isn't a repair. It moves values into compilation and can create one executable per value. Reserve static arguments for true configuration such as activation choice or a small fixed mode set.

Python side effects also happen during tracing. Appending metrics to a global list, incrementing a Python counter, or reading a changing file from a jitted function doesn't describe repeatable device computation. Return metrics and perform I/O in the outer loop.

Measure completed work, not dispatch

JAX dispatches array work asynchronously. Python can regain control before an accelerator finishes. A timer around only result = compiled_fn(x) may measure enqueue overhead rather than completed computation.^{[9]Reference 9Benchmarking JAX codehttps://docs.jax.dev/en/latest/benchmarking.html}

Use .block_until_ready() on a result when timing device execution. Also place inputs on target device before steady-state measurement, match dtypes, and separate transfer time when transfers are part of the real workload.

The benchmark skeleton below distinguishes cold and warm calls. It isn't marked runnable because timing assertions depend on machine load and backend.

Cold time includes compilation and execution. Warm time reuses a compatible executable. Report both when research iteration latency matters. For end-to-end systems, include data loading, host-to-device transfer, evaluation, and checkpoint work rather than quoting only a kernel microbenchmark.

PyTorch accelerator work can also be asynchronous. Keep the general habit: synchronize according to framework and backend before comparing completed work. Framework names don't make a benchmark fair; matched workloads and boundaries do.

Failure cases that look like model bugs

JAX contract violations often produce model symptoms: repeated keys make dropout masks identical, shape changes create compilation pauses that look like input stalls, and captured globals retain old coefficients after configuration changes.

Use symptom, cause, and repair together:

Debug outside in

Remove jit temporarily and run the smallest failing function eagerly. Eager execution won't reproduce every compilation issue, but it gives clearer Python stack traces and confirms base numerical behavior.

Then inspect jax.make_jaxpr(function)(example_inputs) or a lowered representation when you need to see what JAX traced. Compare tree structures and leaf shapes before reading compiler output. Most failures happen at the Python-to-array boundary, not inside generated machine code.

Reintroduce transformations one at a time: base function, then grad, then vmap, then jit. The first boundary that fails names the contract you need to fix.

Preserve JAX when its structure carries evidence

A paper repository contains more than equations. Its tests, checkpoint layout, data order, precision settings, sharding rules, and evaluation scripts define the reported experiment. Rewriting all of that changes many variables at once.

Modify JAX in place when:

• reference checkpoints and evaluation code already use its pytrees;
• one local loss, update, sampler, or transformation needs adjustment;
• vmap, higher-order differentiation, or vectorized simulation expresses core research logic;
• TPU or multi-device infrastructure is already proven;
• paper reproduction requires matching original execution path;
• maintainers review and test JAX changes directly.

PyTorch may remain better for your production stack, team tooling, or supported operators. That doesn't make a research rewrite free. Keep the research implementation as oracle, add cross-framework fixtures, and port only the smallest boundary whose ownership is clear.

Use this decision table before replacing a working codebase:

A clean port requires more than close loss curves. Compare preprocessing, parameter layouts, random seeds and streams, numerical precision, optimizer details, batch ordering, masks, evaluation aggregation, and generated artifacts. A mismatch in any one can dominate framework differences.

A focused modification workflow

Start by running one existing repository test and one tiny forward pass. Save inputs and outputs as a fixture if repository policy allows it. Record tree paths, leaf shapes, dtypes, device placement, and random-key ownership around intended edit.

Make change in smallest pure function. Run it eagerly on fixture. Compare scalar outputs and gradients, then apply same transformations used by repository. Only after correctness passes should you benchmark compiled path.

For a new regularizer, workflow could be:

1. add term to scalar loss function;
2. verify term is zero under known fixture;
3. compare finite differences on tiny float input when useful;
4. check gradient pytree structure and non-finite values;
5. run jitted step twice with same shapes;
6. verify one controlled key change affects only stochastic lane;
7. run held-out training or evaluation slice;
8. measure cold compile and warm execution separately.

That sequence preserves causal evidence. If final metric changes, you can point to one numerical edit rather than a framework rewrite plus a new data loader plus a new optimizer plus a new checkpoint converter.

Keep the execution contract visible

Trace one batch through a JAX repository and name every state boundary. Parameters and optimizer state arrive as pytrees. Randomness arrives as keys. Array updates return values. Transformations derive gradient, batch, and compiled functions from a small numerical core.

Reject misleading diagnoses by checking contract first. Compile pauses may indicate shape churn, repeated dropout masks can expose key reuse, and stale coefficients can come from captured globals. Inspect shape, static values, key ownership, side effects, and synchronization before blaming model math.

Keep these checks close during research edits:

• inspect function inputs and returned state;
• compare pytree structures before leaf values;
• split each random subkey once;
• retain every immutable update result;
• stabilize hot-path shapes and static options;
• add transformations one boundary at a time;
• block before timing completed device work;
• preserve reference fixtures before considering a port.