Frameworks

How the array libraries actually work under the API: tracing a Python function into a typed intermediate form, the transforms (jit, grad, vmap) that rewrite it, the module systems that hold parameters, and the pytree contract that lets any of it compose.

Explainers

• Tracing → jaxpr: the One Trick Behind Every JAX Transform — jit, grad, and vmap look like magic, but they all do the same mundane thing first: they run your Python once, feeding it stand-in values that carry only a shape and a dtype, and record every operation into a small typed program called a jaxpr. Watch a real function become its jaxpr line by line, see the print that fires exactly once, learn why a new input shape makes jit retrace, why a plain Python if throws under jit, and how the pytree flatten/unflatten contract lets all of it compose. Every other JAX page rests on this one.
• jit(grad(vmap(f))): Why Transform Order Changes the Answer — jit, grad, and vmap all take a function and return a function, which makes them feel like decorators you can stack in any order. They are not: vmap(grad(f)) computes one gradient per example — the standard trick behind per-example gradients — while grad(vmap(f)) does not even type-check until you bolt on a reduction, and which reduction you pick changes the number you get. We ground it in JAX’s own tracer classes, reorder a live jit/grad/vmap chip stack and watch the trace-stack diagram and a real numeric example recompute together, then find the same order-sensitivity again in jacfwd vs jacrev.
• Params Are Data: the Functional Model Behind Flax — A Flax model never holds its own weights. Call init() and it hands back an ordinary nested dict of arrays; call apply() and it takes that same dict back in to run the forward pass. We follow one PRNG key as it splits down the module tree, watch every self.param drop a leaf into the growing params pytree, send that pytree back through apply, and catch the exact bug Flax’s own docs name when submodules get built inside a branch — the discipline that lets checkpointing, sharding, and jax.grad treat a whole model like any other array.