How JAX and PyTorch run your code

One function, two routes

Take a function with three array operations: a matrix multiply, an elementwise tanh, and a sum.

@jax.jit
def step(w, x):
    h = jnp.tanh(x @ w)
    return h.sum()

step(w, x)   # first call; w and x are f32[512, 512]
step(w, x)   # second call, same shapes

The PyTorch version drops the decorator and calls torch.tanh. When step is called, jax.jit first looks for a compiled executable that matches the shapes and dtypes of the arguments. On the first call there is none, so JAX runs the Python body once with stand-in values, records the operations, and hands the record to the XLA compiler. The second call finds the executable and runs it. Eager PyTorch keeps no compiled executable for this function: every call runs the Python body and dispatches tensor operations as Python reaches them. An operation can launch one or more GPU kernels.Dispatch is asynchronous in both frameworks. Python queues device work and moves on; it waits only when it needs a result on the host, for example to print it. In JAX, block_until_ready() waits explicitly; in PyTorch, torch.cuda.synchronize().

The recorded program has one equation per operation: a dot_general for the matrix multiply, a tanh and a reduce_sum. XLA compiles the three as a whole, so it can fuse the elementwise work into fewer kernels than there are operations.Whether XLA fuses a given pair depends on the backend and the shapes. The executable contains one or more device kernels.

Compilation is an explicit choice in both frameworks. Without jit, JAX dispatches operations one at a time; torch.compile captures PyTorch code into graphs and compiles them. The comparison here is between execution modes as well as frameworks.The JAX JIT guide contrasts an undecorated function with its jax.jit version. Applying grad or vmap does not by itself compile the whole function.

The layers between Python and the chip

Both stacks have the same kinds of layer: an API, an intermediate form, a compiler, kernels, a runtime and hardware. Eager PyTorch goes from the API to kernels directly. A call such as torch.matmul selects an implementation for the tensor’s device and dtype and queues it on the accelerator.

jax.jit puts two layers in between. It traces the function into a jaxpr, lowers the jaxpr to StableHLO, and XLA compiles StableHLO into an executable for the target device.StableHLO is the MLIR operation set that XLA accepts as input. torch.compile adds a similar pair of layers to PyTorch: TorchDynamo captures graphs, and TorchInductor compiles them.

How jit traces a function

jax.jit does not read your source code. It calls your function with tracers, stand-in objects that carry a shape and a dtype but no values, and records every JAX operation they pass through.

The record is a jaxpr, a small functional language with one equation per primitive operation. jax.make_jaxpr(f)(x) prints it.

Ordinary Python runs during tracing and not afterward. A print inside a jitted function runs on a cache miss and prints a tracer. On a cache hit, JAX runs the compiled executable and the print does not run. jax.debug.print prints from inside the compiled program instead.

A Python if cannot branch on a traced value. if x.sum() > 0: asks the tracer for a Python bool, and the tracer has no value to give, so tracing fails with TracerBoolConversionError. lax.cond records both branches and leaves the choice to the compiled program.

Compiled executables are specialized to their inputs. JAX caches them by abstract signature: shapes, dtypes and the values of static arguments. A call with a new shape traces and compiles again.

Eager PyTorch runs the Python function on every call, so data-dependent control flow works directly, though turning a CUDA tensor into a Python bool makes the host wait for the GPU. With torch.compile, TorchDynamo captures graph regions and checks guards before reusing them. Code it cannot capture causes a graph break: PyTorch compiles the region captured so far, runs the unsupported code in Python, and resumes capture after it.

Fusion reduces memory traffic

An NVIDIA H100 SXM performs about 989 trillion dense bf16 operations per second and moves about 3.35 TB per second to and from its HBM memory.Vendor peaks from NVIDIA’s H100 specifications, rounded and approximate. NVIDIA also quotes “with sparsity” figures that are twice the arithmetic rate. That is roughly 295 operations for every byte moved. Adding two float32 arrays does one operation per element for 12 bytes of traffic: two 4-byte reads and one 4-byte write, so its speed is set by memory traffic, not arithmetic.

For the chain modeled below, eager execution launches a separate kernel for each operation and materializes intermediate arrays in device memory. Caches can serve some reads; the model counts transfers to and from HBM. A compiler that sees the whole chain can fuse compatible operations into one kernel and keep the intermediate values in registers or on-chip memory. XLA does this under jax.jit, and TorchInductor, the default backend of torch.compile, does it for PyTorch. Reductions such as softmax and library calls such as matrix multiplication often end a fused group, although compilers can fold pointwise work into a matrix multiply’s output stage.

In this peak-rate model the matrix multiply is limited by arithmetic. Its 2 × 4096³ ≈ 137 billion operations take about 0.28 ms at 495 TFLOP/s, the H100’s dense TF32 peak, while its 192 MiB of traffic needs about 0.06 ms.The estimate assumes the float32 matrix multiply runs on TF32 tensor cores at the vendor peak; real kernels reach less. What fusion removes is the memory traffic and launch overhead of the small operations around it.

Two ways to compute a gradient

Both frameworks implement reverse-mode automatic differentiation. The function runs forward; then the chain rule runs backward from the output, multiplying one local derivative per operation. The frameworks differ in when the backward pass is built and where its result goes.

PyTorch records a graph while the forward pass runs. Each operation on a tensor that requires a gradient appends a node such as SinBackward0, which keeps whatever its derivative will need. loss.backward() walks the graph in reverse and adds each gradient into the .grad attribute of a leaf tensor such as a parameter. The graph holds exactly the operations that ran, and it is freed after the backward pass.

jax.grad(f) returns a new function. JAX linearizes the traced computation and transposes its linear part, which produces the reverse pass as ordinary jaxpr equations. Calling the new function returns the gradient as a value, so it composes with jit, vmap and further differentiation.

Transformations take a function and return a function

jit, grad and vmap each take a function and return a new one, so they nest. vmap adds a batch dimension. It applies a batching rule to each primitive, so code written for one example accepts a whole batch without a Python loop. In figure 6 a dot product of two vectors becomes one dot_general between a vector and a matrix.

The order of nesting decides the program. vmap(grad(loss)) differentiates the loss for one example and then batches that, which gives per-example gradients. grad(vmap(loss)) fails, because vmap(loss) returns a vector and grad needs a scalar output. PyTorch provides the same transformations in torch.func, including torch.func.grad, torch.func.vmap and torch.func.jacfwd, and torch.compile can compile compatible regions of the result.

Where state lives

Write code under JAX transformations as pure functions: return results explicitly instead of changing hidden state. Parameters, optimizer state and random keys enter as arguments and leave as return values, and updating an array produces a new array. XLA may still reuse a buffer in place when its analysis shows the result is the same.

Typical PyTorch training code keeps parameters inside an nn.Module and updates them through an optimizer, and random draws come from a default generator kept for each device. That style is shorter to write. The cost is that the function signature does not list everything the step reads or changes.

Randomness shows the difference most clearly. A JAX sample is a deterministic function of its key, and jax.random.split derives new keys from an old one. A PyTorch draw advances a hidden generator, so its value depends on how many draws came before it. Figure 9 adds one debugging draw at the start of a step that already makes three.Keep the key assigned to each draw fixed. For example, use jax.random.fold_in(root, id) with distinct, stable integer IDs for mask, noise, initialization and debugging. Adding an ID leaves the existing keys unchanged; changing the key assignments can change the draws. See the fold_in API.

This is a difference in how the examples manage state. PyTorch can also isolate debugging draws with a separate torch.Generator. In JAX, using the same key twice repeats the sample rather than advancing a stream.

How TPUs and GPUs schedule a matrix multiply

A TPU TensorCore contains matrix-multiply units (MXUs) alongside vector and scalar units. An MXU is a systolic array: a grid of multiply-accumulate cells through which operands move one cell per clock cycle, on a schedule the compiler fixes before the program runs. MXUs are 128 × 128 cells before TPU v6e and 256 × 256 in v6e and TPU7x, so one MXU performs 16,384 or 65,536 multiply-accumulates per cycle when full.For chip-level figures such as HBM capacity, bandwidth and pod sizes, see the Google Cloud accelerators explainer in this collection.

An NVIDIA GPU divides the same work into thread blocks, groups of threads that run in warps of 32. A hardware scheduler on the chip assigns blocks to streaming multiprocessors (SMs) as they become free, while the kernel runs. An H100 SXM has 132 SMs, each with tensor cores for matrix multiplication, registers and shared memory. Irregular workloads, dynamic shapes and hand-written kernels fit this model, because the hardware finds parallelism at run time.

Sharding a computation across devices

JAX can lay out an array across a named mesh of devices. A PartitionSpec says which array dimension is split along which mesh axis. Under jit, the compiler propagates those layouts through the program and inserts collectives, operations that move data between devices, where values must move: an all-gather to reassemble a split array, an all-reduce to sum partial results, a reduce-scatter to sum and split again.

PyTorch offers the same patterns at several levels. DistributedDataParallel replicates the model and all-reduces gradients. FullyShardedDataParallel shards parameters and gathers them before use. DTensor describes tensor placements over a device mesh, close to the JAX model.

Choosing between them

The choice follows from the execution model you want, the hardware you have, the libraries you depend on, and how you like to debug.

JAX fits when

• You train or serve on TPUs, or scale one program across hundreds of chips, where compiler-managed sharding saves work.
• Your shapes are stable and the workload gains from optimization across large regions of the program.
• You need higher-order derivatives, per-example gradients or other function transformations as everyday tools.

PyTorch fits when

• You want eager execution by default while you iterate on dynamic models or data-dependent control flow.
• Your dependencies and deployment tools are built around the PyTorch ecosystem.
• Your team prefers stateful modules and wants to add compilation region by region.

Further reading

• How to think in JAX and JAX: the sharp bits, from the JAX documentation.
• Autodidax, which rebuilds the core of JAX, tracing included, in one notebook.
• How to Scale Your Model, on TPU architecture and sharding arithmetic, from Google DeepMind.
• OpenXLA for the compiler, and the torch.compile documentation for PyTorch’s compiler path.