Coming from PyTorch, the first thing JAX does is offend you. You write a simple function with an if statement, wrap it in jax.jit, and call it. Instead of running, it throws a TracerBoolConversionError. No graceful fallback. No partial compile. Just a hard stop on a line of Python that would have run fine a second ago.
torch.compile would have shrugged, taken a graph break, and moved on. JAX refuses.
That refusal is the most important thing to understand about JAX. Once you accept that JAX is not a framework with a compiler attached, but a tracing machine, the whole system falls into place. Every headline feature (jit, grad, vmap, shard_map) is the same move applied repeatedly: run your pure Python function once with stand-in values, record the sequence of primitive operations into a small, typed IR called a jaxpr, then transform or compile that IR.
Understanding JAX means understanding the trace boundary: what the machine can see when your function runs, and what it is structurally blind to.
The jaxpr is the whole point
Start with the object everything revolves around. jax.make_jaxpr traces a function and hands you the IR without compiling it:
import jax.numpy as jnp
from jax import make_jaxpr
def f(x):
return jnp.maximum(x * 2.0 + 1.0, 0.0) # relu(2x + 1)
print(make_jaxpr(f)(jnp.float32(3.0)))
# { lambda ; a:f32[]. let
# b:f32[] = mul a 2.0:f32[]
# c:f32[] = add b 1.0:f32[]
# d:f32[] = max c 0.0:f32[]
# in (d,) }Read it like a typed lambda in A-normal form. After the semicolon come the input binders, here one scalar, a:f32[]. Then a let block of primitive applications, each binding a typed intermediate, and finally the returned tuple. Every value carries its type as a dtype and a shape.
Notice what is not there. The number 3.0 you passed in is gone. The jaxpr is not a trace of one execution; it is the abstract computation, parameterized over the types of its inputs, with all concrete data erased.
Why JAX refuses data-dependent control flow
The refusal is structural. A jaxpr is a static graph: every branch must be known at trace time. A data-dependent if cannot be represented, because the machine cannot know which way it goes without running it, and running it would defeat the purpose of tracing. The escape hatches are jax.lax.cond, jax.lax.while_loop, and jax.lax.scan, which encode both branches into the graph and let the compiler handle the dispatch.
This is the sharpest contrast with torch.compile, which takes graph breaks and falls back to eager. JAX's bet is that a hard refusal at trace time is better than a slow path at runtime. The compiler sees the whole program or it sees nothing.
Randomness becomes a pure function
JAX's PRNG is a pure function of a key: the same key produces the same sequence, deterministically. There is no hidden global state. This is what makes jit, grad, and vmap all work with randomness: the key is just another argument, threaded through the computation like any other value.
The abstract value has quietly evolved
The abstract value (aval) has quietly evolved to include sharding and memory space. What started as dtype and shape now carries where the data lives and how it is laid out. This is the change that made placement part of the type rather than an annotation, and it is what makes JAX's sharding model work.
JAX is a photographer who insists on taking the photo before you move. torch.compile is a videographer who films everything and edits later. JAX's way is stricter, but the final image is sharper because nothing was blurred by motion.
JAX is not a framework with a compiler attached. It is a tracing machine.