JAX for PyTorch Researchers

You change a loss term in a paper's JAX repository. The numbers look right, but a short final batch triggers another compilation and dropout repeats the same mask. Those failures can come from shapes and random-key ownership rather than the model's mathematics.

You already know tensors, gradients, and training loops from PyTorch. We'll keep a four-row linear-regression example small enough to check by hand while making parameters, updates, and randomness explicit. The same checks help you modify a larger research implementation without replacing its framework.

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.

Start by tracing one call, not by translating a whole framework. Find the batch entering the loss, the values returned by one step, and the point where a transformed function crosses the Python boundary. Once you can name those edges, later choices about vmap, jit, or sharding have somewhere concrete to attach.

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.

While tracing that call, identify changing state, shape/static arguments, random-key ownership, and transformed functions. For our example, the model predicts weight[0] * x[0] + weight[1] * x[1] + bias. Its four rows are:

The initial mean squared error is (4 + 1 + 1 + 9) / 4 = 3.75. We know an exact solution too: weights [2, -1] and bias 0. Those numbers give us an independent check before compilation enters the story.

That inventory often reveals a narrow edit. Put a new regularizer in the loss, change sampling where a one-use key enters, or derive a per-example statistic with vmap; leave the surrounding research harness alone until evidence says it's the problem.

Neither framework has only one programming style. The table is a reading lens, not a law.

One step is a state transition

In the GPT lab, mutation was the whole loop: model.train(), optimizer.zero_grad(), loss.backward(), optimizer.step(). Parameters changed on the module, gradient buffers appeared on tensors, and a random generator advanced behind dropout or sampling. The call site can observe the result, but its return type doesn't tell you which pieces changed.

With Adam and dropout, the caller must retain updated parameters, updated optimizer state, a continuation key, and metrics. Stateful layers may add statistics such as a running mean. JAX transformations work cleanly when those changing values become ordinary inputs and outputs.^{[2]Reference 2Stateful computationshttps://docs.jax.dev/en/latest/stateful-computations.html}

Read the figure as that caller's ledger. Every arrow marks ownership: the step consumes current state and hands back the next state.

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}

Predict what happens if a jitted loss reads weight_decay from a module-level variable. The first trace can capture 0.01; changing the Python global to 0.1 doesn't guarantee the cached executable sees that edit. Pass research settings as arguments, then decide whether each one is an array value or a small static configuration.

Printing inside a jitted function exposes the same boundary. Python print runs while JAX traces, so it may display a tracer once instead of a device value on every execution. Return metrics, use jax.debug.print when you need a transformed-code debug print, or inspect values outside the compiled boundary.

Immutable arrays remove aliasing surprises

In NumPy and PyTorch, writing x[i] = value mutates the underlying memory buffer in place. That feels convenient in eager Python scripts, but in an optimizing compiler such as XLA, in-place mutation turns optimization into an alias-analysis nightmare: the compiler can't easily prove whether another concurrent thread or device stream is reading that same memory address. JAX enforces pure value semantics: ordinary jax.Array instances are immutable. Attempting x[i] = value raises TypeError immediately.^{[3]Reference 3JAX: The Sharp Bitshttps://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html}

To update an element, JAX provides the .at[...] indexing operator: updated = x.at[i].set(value). This expression returns a fresh array value representing the updated state; the original x remains completely untouched.

You might worry that returning a new array forces a full memory copy on every slice update. Under jax.jit, XLA analyzes buffer lifecycles. If the compiler proves the old buffer has no other live references (or if you explicitly donate the buffer via donate_argnums), XLA reuses the physical memory buffer in place without an extra allocation. You get the clean mental model of functional immutability in Python alongside the zero-copy speed of mutating memory on accelerator hardware. Controlled mutable Ref APIs exist for low-level kernel plumbing, but standard research training loops operate purely on immutable arrays.^{[4]Reference 4Ref: mutable arrays for data plumbing and memory controlhttps://docs.jax.dev/en/latest/array_refs.html}

The examples were checked with JAX/jaxlib 0.11.1 on one CPU device using Python 3.12. Install the matching jax==0.11.1 package for these fixtures; accelerator setup is outside this lesson.

After updated = original.at[1].set(99), should original show 20 or 99? Compare both bindings in this example.

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.

Arrays aren't the only values that used to hide behind mutation. Randomness does too.

Randomness is explicit data

In PyTorch, calling torch.randn(...) reads and advances a hidden global generator state. That model breaks down on accelerator clusters: thousands of concurrent threads competing to update a single sequential state register cause lock contention, race conditions, or device-order nondeterminism. Worse, an optimizing compiler can't reorder or fuse operations without accidentally scrambling the random sequence.

JAX eliminates global RNG state entirely by using counter-based pseudorandom number generation, specifically the Threefry algorithm. In a counter-based PRNG, generating random numbers isn't a stateful mutation: it's a pure cryptographic permutation mapping an explicit key and counter to bits: $\text{bits} = \text{Threefry}(K, C)$. Random generation becomes embarrassingly parallel across devices and cores, with zero synchronization overhead.

The golden rule to master: JAX random functions are pure mathematical functions. Passing the same key twice yields the exact same bits. Sampling doesn't advance or mutate the key. To draw independent samples, you explicitly split a key into one-use child keys and a continuation key using jax.random.split(key, num).^{[5]Reference 5Pseudorandom numbershttps://docs.jax.dev/en/latest/random-numbers.html}

Split for independent operations, 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 makes that ownership rule visible: key reuse repeats a sample, while separate subkeys produce different samples.

Typed keys from jax.random.key() carry their PRNG implementation in the dtype. The printed key<fry> is a scalar typed key for Threefry, not a length-two numeric array, although its underlying key data consists of two uint32 words. Older repositories may use jax.random.PRNGKey(), which returns that legacy numeric representation. Read the key format before changing checkpoint serialization; random.key_data and random.wrap_key_data provide explicit conversions.^{[6]Reference 6jax.random modulehttps://docs.jax.dev/en/latest/jax.random.html}

JAX doesn't promise sequential equivalence: drawing three scalars from split keys needn't give the same sequence as random.normal(key, (3,)) from the parent. This design leaves room for vectorization and parallel execution. To reproduce a sampler, preserve its key derivation, draw shapes, implementation, and software version, not only its seed.^{[5]Reference 5Pseudorandom numbershttps://docs.jax.dev/en/latest/random-numbers.html}

The comparison starts both programs from the same parent seed on purpose. Don't copy that reuse into a training step.

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.

JAX also has an experimental stateful RNG backed by Ref objects. Its current documentation recommends explicit keys for performance-critical work and excludes returning the wrapper through transformations, carrying it through scan, or using it with remat. We use explicit keys throughout; this CPU lab doesn't exercise that experimental wrapper.^{[7]Reference 7jax.experimental.random.stateful_rnghttps://docs.jax.dev/en/latest/_autosummary/jax.experimental.random.stateful_rng.html}

A dropout helper receives one key and applies two independent dropout masks. It calls random.bernoulli(key, ...) for both masks. Tests are reproducible, but both masks always match. What contract is broken, and where should you repair it?

Answer

The helper reused one key for two independent random events. Split its input key once at the helper's ownership boundary, consume one subkey per mask, and return or retain a separate continuation key if the caller needs future randomness. Changing the seed would only produce a different pair of matching masks.

Once keys and arrays are ordinary data, transformations can wrap the function without guessing hidden state. The next question is which artifact each wrapper should produce.

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, and check its output eagerly. Then decide what evidence you need: a derivative, one result per example, or a reusable compiled call. Add grad, vmap, or jit for that artifact, one boundary at a time.

The flow below shows function derivation, not runtime stages. Each box is still 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.^{[8]Reference 8Automatic 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.^{[9]Reference 9Automatic 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.^{[10]Reference 10Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

Transformation order changes the derived program

One composition preserves a gradient for each example; the other differentiates one aggregate loss. These expressions are sketches of the two derived functions, with complete definitions in the runnable cell below:

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.

For the four-row regression example, the first weight's per-example gradients at zero are [-4, 0, -2, -12]. Their mean is -4.5. The second weight's mean gradient is -1.5, and the bias gradient is -2.5. One SGD step with learning rate 0.05 should therefore produce weights [0.225, 0.075] and bias 0.125.

Derive those quantities with JAX and check that averaging per-example gradients matches differentiating the mean. These linked cells share state; run them in order. They use CPU-sized float32 arrays and require JAX, with no accelerator installation assumed.

The extra axis belongs to examples: the weight-gradient leaf is [4, 2], whereas the mean-gradient leaf is [2]. A length-one array isn't a scalar either; grad expects an output with shape () for this ordinary real-valued use.

Gradients need a place to live. That place is the same nested structure as the parameters.

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.^{[11]Reference 11Pytreeshttps://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.

Renaming one key from w to kernel can leave the leaf count unchanged while changing the state schema. A correct update must match paths and shapes, not merely the number of arrays.

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

Two trees passed to tree.map need compatible structure. But matching structure doesn't validate leaf shapes: ordinary array arithmetic may silently broadcast an incorrect gradient. For parameter updates, check both the tree definition and exact leaf shapes before applying a leafwise update.

For example, a parameter tree keyed by {"layer": {"w": ..., "b": ...}} can't safely update from a gradient tree keyed by {"layer": {"kernel": ..., "b": ...}}. The leaf shapes may match, but the renamed key changes which state the update refers to.

Exercise both failure modes. A renamed key is a tree error. A scalar gradient under the correct key can broadcast without an error, so the explicit shape check below rejects it before an update.

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 the same element count and still changes model behavior.

You need per-example gradient norms for clipping, but the current step computes the gradient of mean batch loss. Should you wrap the entire compiled training step in vmap, or derive a narrower function?

Answer

Derive grad(loss_one) from the scalar one-example loss, then vmap that gradient over example axes while sharing the parameter pytree with in_axes=None. Wrapping the whole stateful step would batch optimizer state and updates too, materialize more data than needed, and obscure whether aggregation happens before or after differentiation.

A training step is those four pieces: a pure loss, a gradient pytree, an explicit update, and optional vmap. The GPT lab's loss.backward(); optimizer.step() is the same state transition. The smallest version you can run doesn't need a Transformer.

Translate a tiny training step

The familiar PyTorch loop makes mutation explicit through method calls. This abbreviated comparison assumes an existing model, optimizer, and tensors; it isn't a standalone program.

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.

Continue with the same four-row model. The code checks the first update against our hand calculation, trains for 80 steps, and evaluates the returned parameters. The loss returned by train_step is measured before that step's update, so compute a fresh final loss rather than relabeling the last returned loss.

The update has no optimizer state because plain stochastic gradient descent (SGD) 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.

Notice what the example makes easy to test: one scalar loss, one parameter tree, one returned update, and one reported value. The same boundary survives when an optimizer adds state or a model adds dropout; only the tuple of explicit inputs and outputs grows.

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.

Add two real pieces of state to the same regression model: momentum and input dropout. Momentum retains a moving update direction, new_velocity = 0.9 * velocity + gradient. Input dropout retains each feature with probability 0.75 and divides retained values by 0.75 so their expectation is unchanged. This is a state-ownership exercise, not a claim that dropout improves this tiny regression problem.

Replay identical inputs to check deterministic behavior. Then thread the returned state into a second step. A different continuation key is the invariant; two small dropout masks could coincide by chance.

The replay deliberately reuses an input key to test reproducibility. A real training loop keeps key1, then key2, rather than repeatedly submitting the original key. Resuming this step requires parameters, velocity, and the continuation key, together with the data-loader position.

The step works. Now ask whether the next call has the same shape, dtype, static settings, and device placement. Compilation is the next contract, because the first slow call isn't the same as the hundredth.

JIT compiles a shape-specific contract

Just-in-time (JIT) compilation has an up-front cost. When you call a @jax.jit-decorated function for the first time, JAX executes a two-stage compilation pipeline:

1. Abstract Tracing: JAX runs your Python function once using abstract tracers (ShapedArray). These tracers track array shapes and dtypes without concrete numbers, recording every primitive operation into an intermediate representation graph called a Jaxpr (JAX expression).
2. XLA Lowering & Compilation: JAX lowers the Jaxpr to StableHLO, and the XLA compiler generates an optimized machine-code binary for your accelerator. Kernel fusion, buffer reuse planning, and dead-code elimination happen here.

XLA caches the compiled executable in memory using a cache key formed by four attributes: the function's object identity, input array shapes, input dtypes, and the exact values of arguments marked static. Subsequent calls matching this exact signature bypass Python execution completely and run the cached binary directly on hardware.^{[10]Reference 10Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

A [32, 512] token batch and a [16, 512] final batch don't share the same compiled variant. The shapes differ, producing a cache miss that forces a brand-new compilation. Float32 versus bfloat16 inputs split variants too, and changing a static argument recompiles by definition. A sudden multi-second pause at the end of an epoch is almost always compilation doing exactly what your changed batch shape demanded.

The short remainder isn't necessarily a data bug. It's a new input contract, so JAX may trace and compile again. If that remainder is expected, pad to a fixed shape and carry a mask, bucket sequence lengths into a small set, or handle the short batch on a deliberate side path.

Keep static option values bounded and construct the jitted callable once outside hot loops. Log compilation separately from execution so an intentional variant is visible rather than mistaken for kernel slowness.

Expose a captured-global bug alongside a shape change. The list below intentionally records Python tracing events for diagnosis; it isn't a device-side metric counter. A new shape retraces and captures the new global, so the bug can look intermittent.

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.^{[10]Reference 10Just-in-time compilationhttps://docs.jax.dev/en/latest/jit-compilation.html}

Before choosing a fix, decide whether the branch belongs to data or configuration. A per-example mask is data and must stay in the array program; an activation choice from a small fixed set can be static and may justify bounded recompiles.

Choose a fix by semantics:

Both branches of lax.cond are traced. A scalar condition selects one branch at execution, subject to compiler transformations; vmap over predicates can turn it into elementwise selection.^{[12]Reference 12jax.lax.condhttps://docs.jax.dev/en/latest/_autosummary/jax.lax.cond.html} jnp.where isn't a lazy Python if: don't rely on it to make invalid arithmetic or NaN gradients in an unselected expression harmless.

Loop carries must retain fixed shapes and dtypes. Also, lax.while_loop doesn't support ordinary reverse-mode differentiation. A fixed-length lax.scan is often the appropriate training-loop form when you need reverse-mode gradients through iterations.^{[13]Reference 13jax.lax.while_loophttps://docs.jax.dev/en/latest/_autosummary/jax.lax.while_loop.html}

Run both sides of a repaired scalar branch, then expose the reverse-mode limitation instead of assuming that every jittable program is differentiable.

The fixed three-step loop has a different contract from "double until at least eight." The example compares their differentiation support; it doesn't claim the two functions agree for every input.

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.

Compilation and dispatch also need separate measurements.

Measure completed work, not dispatch

JAX dispatches array work asynchronously. Python can regain control before an accelerator finishes. Before trusting a timer, ask what it has observed: the Python enqueue call, or completed device work?

A timer around only result = compiled_fn(x) may measure enqueue overhead rather than execution.^{[14]Reference 14Benchmarking JAX codehttps://docs.jax.dev/en/latest/benchmarking.html}

Use .block_until_ready() on a result when timing device execution. Place inputs on the target device before steady-state measurement, match dtypes, and separate transfer time when transfers are part of the real workload. JAX defaults to 32-bit dtypes, so a "JAX vs NumPy" comparison that leaves NumPy in float64 isn't a framework result.^{[14]Reference 14Benchmarking JAX codehttps://docs.jax.dev/en/latest/benchmarking.html}

An array has a .block_until_ready() method; a tuple or dictionary doesn't. For the momentum step's structured return, jax.block_until_ready(result) waits on its array leaves, as the state example does above.

This runnable CPU microbenchmark uses a fixed [512, 64] float32 input and waits for every timed result. It checks shape and finiteness, not a machine-dependent speed threshold. Its printed timings will vary, so no timing snapshot is embedded here.

Cold time includes tracing, compilation, dispatch, and execution for this fresh callable; a persistent compilation cache can change that breakdown. Warm time includes Python dispatch and a wait per call, so it's a synchronized per-call latency measurement. Enqueuing many independent calls and waiting once instead measures an amortized throughput-oriented boundary, not the same latency. Don't label one as the other.

For accelerator measurements, select and report the actual device, JAX/jaxlib versions, dtype, shapes, compilation-cache policy, transfer boundary, and synchronization rule. This lab verifies CPU numerical behavior and timing structure only. It doesn't establish GPU/TPU speed or multi-host performance. End-to-end measurements must also account for data loading, transfers, evaluation, and checkpoint work when those are in scope.

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.

A benchmark reports JAX is 40 times faster than PyTorch. JAX timer covers one call after compilation but never reads result; PyTorch timer includes first accelerator call and synchronization. What must change before speed claim is meaningful?

Answer

Measure matched boundaries. Put inputs on target devices, synchronize both frameworks, report JAX cold compile-plus-execution separately from warm execution, warm PyTorch path too, and include transfers only when both measurements include them. Without those controls, comparison mixes dispatch latency, compilation, execution, and transfer costs.

Treat devices as another state boundary

The single-device path now has a measurable contract. Remember that vmap is local vectorization on one device: it doesn't distribute work across hardware accelerators.

When scaling across multiple GPUs or TPU pods, JAX provides two architectural generations of distributed execution:

1. Historical SPMD with jax.pmap: Parallel map maps a function across device replicas by expecting a leading device dimension on input arrays (for example, shape (8, 32, 512) across 8 GPUs). Each device runs the identical compiled program on its local shard, synchronizing through collective primitives such as lax.psum.^{[15]Reference 15jax.pmaphttps://docs.jax.dev/en/latest/_autosummary/jax.pmap.html}
2. Modern Array Sharding with shard_map & NamedSharding: Current JAX treats multi-device distributed parallelism as an array property rather than a function wrapper. You define a logical device mesh (jax.sharding.Mesh(devices, ('data', 'model'))) and annotate PyTree leaves with a PartitionSpec. Inside jax.jit, XLA inspects the sharding specs and automatically inserts cross-device communication collectives (all-reduce, reduce-scatter, all-gather). When you need explicit, manual per-device SPMD logic, shard_map lets you write per-device kernels with explicit collectives without leaving jax.jit.

Explicit sharding adds another layer to debug. An array can have the expected global shape while its physical shards sit on an unexpected mesh or trigger silent communication transfers between operations.

Never copy a single PRNG key into every device lane: doing so makes every accelerator draw identical dropout masks, eliminating regularized diversity across data-parallel replicas. Use jax.random.fold_in(key, device_id) or split the key across the device axis so each replica derives an independent, reproducible random stream.

If a research repository already uses pmap or explicit sharding, preserve its version-specific wrappers until you've verified local shard shapes, device placement, numerical parity, and synchronized timing. A fast replicated call with the wrong local batch or duplicated randomness isn't a successful port. No multi-device collective, sharding strategy, or checkpoint migration is executed by this lesson's CPU tests.

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. Start by classifying the symptom, then check the boundary that could produce it.

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; those checks can isolate the boundary without requiring you to interpret 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. If the failure appears only after adding devices, repeat the same process with placement and collective boundaries visible.

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.

Keep the JAX path when the evidence already lives there. Reference checkpoints and evaluation may depend on its pytrees, or one local loss, update, sampler, or transformation may be the only thing that needs adjustment. vmap, higher-order differentiation, and vectorized simulation can express core research logic.

Proven TPU (Tensor Processing Unit) or multi-device infrastructure can make a rewrite especially costly. Paper reproduction may require the original execution path, and maintainers may already 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 the intended edit.

Change the smallest pure function and run it eagerly on the fixture. Compare scalar outputs and gradients, then apply the same transformations the repository uses. Benchmark the compiled path only after those checks pass.

If you're adding a regularizer, treat this sequence as an evidence ladder:

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

That sequence preserves causal evidence. If the 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.

For practice, add 0.1 * sum(weight ** 2) to the regression loss without regularizing bias. At zero weights, both the extra loss and its gradient should be zero. At weights [2, -1], the extra loss should be 0.5 and its gradient [0.4, -0.2]. Check those values eagerly and under jit, then rerun the training fixture. This tests the intended numerical edit independently of its effect on convergence.