JAX / PyTorch

Everything starts as framework code: jnp operations on whole arrays, module trees, Python control flow. The essential fact about this layer is that it is not what runs. When you call jax.jit, JAX executes your Python once with tracer objects instead of numbers, records every primitive operation the tracers pass through, and keeps the recording. Your function becomes data.

Tracing explains most early JAX surprises. Python control flow that branches on values disappears into whichever branch the trace took. Shapes are fixed at trace time, which is why a new shape triggers a recompile. Print statements run once, at trace time, not per step. None of these are quirks; they are what "your function becomes data" means in practice.

PyTorch reaches this same stack through a different door: TorchTPU and torchax map ATen operations into StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, two layers below. By the time either framework reaches the compiler, the framework identity is gone. That shared convergence point is why everything you learn descending this stack applies to both.

def attention(q, k, v):
    s = q @ k.T
    m = jnp.max(s, axis=-1, keepdims=True)
    p = jnp.exp(s - m)
    l = jnp.sum(p, axis=-1, keepdims=True)
    w = p / l
    return w @ v

Carry one question downward: which of these Python lines will still be visible three layers below, and in what form? The matmuls survive. The softmax dissolves into five primitives. The broadcasting that Python hid becomes explicit. The next layer, the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, is where you first see all of that written down.