{"slug": "what-jax-traces-and-what-it-refuses", "title": "What JAX Traces, and What It Refuses", "summary": "JAX, a numerical computing library, is fundamentally a tracing machine that converts pure Python functions into a typed intermediate representation called a jaxpr, refusing to execute code with control flow that cannot be traced, unlike PyTorch's torch.compile which falls back to a graph break. The jaxpr, which erases concrete data and preserves only types and shapes, is the core object that all JAX transformations—jit, grad, vmap, shard_map—operate on, with the library exposing 285 primitive operations as of version 0.11.0.", "body_md": "# What JAX Traces, and What It Refuses\n\nComing from PyTorch, the first thing JAX does is offend you. You write an honest little function with an `if`\n\nin it, wrap it in `jax.jit`\n\n, call it, and instead of running it throws `TracerBoolConversionError`\n\n. No graceful fallback, no partial compile — a hard stop on a line of Python that would have run fine a second ago. `torch.compile`\n\nwould have shrugged, taken a graph break, and moved on. JAX refuses.\n\nThat refusal is the most important thing to understand about JAX, and once it clicks the whole system falls into place. JAX is not a framework with a compiler attached. It is a **tracing machine**. Every headline feature — `jit`\n\n, `grad`\n\n, `vmap`\n\n, `shard_map`\n\n— is the same move applied over and over: run a pure Python function once with stand-in values, record the sequence of primitive operations into a small typed IR called a *jaxpr*, and 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.\n\nEverything below runs on JAX 0.11.0, CPU backend, and every jaxpr and every number is copied from a session rather than written from memory. 1 Line numbers refer to that release. This is the front end: how Python becomes a jaxpr, why control flow is special, why every transformation is a rewrite of the same object, and — the part that has changed most in the last year — what the type of a traced value has quietly grown to include.\n\n## The jaxpr is the whole point\n\nStart with the object everything revolves around. `jax.make_jaxpr`\n\ntraces a function and hands you the IR without compiling it:\n\n```\n1\n2\n3\n4\n5\n6\n7\npython\nimport jax.numpy as jnp\nfrom jax import make_jaxpr\n\ndef f(x):\n    return jnp.maximum(x * 2.0 + 1.0, 0.0)   # relu(2x + 1)\n\nprint(make_jaxpr(f)(jnp.float32(3.0)))\n1\n2\n3\n4\n5\njs\n{ lambda ; a:f32[]. let\n    b:f32[] = mul a 2.0:f32[]\n    c:f32[] = add b 1.0:f32[]\n    d:f32[] = max c 0.0:f32[]\n  in (d,) }\n```\n\nRead it like a typed lambda in A-normal form. 2 After the semicolon come the input binders — here one scalar,\n\n`a:f32[]`\n\n. Then a `let`\n\nblock of primitive applications, each binding a typed intermediate, and finally the returned tuple. Every value carries its type as a dtype and a shape: `f32[]`\n\nis a float32 scalar, `f32[4,8]`\n\na matrix. Notice what is *not*there. The number\n\n`3.0`\n\nyou 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.Give it real arrays and the primitives get more interesting:\n\n```\n1\n2\n3\n4\n5\npython\ndef layer(x, W, b):\n    return jnp.maximum(x @ W + b, 0.0)\n\nx = jnp.ones((4, 8)); W = jnp.ones((8, 16)); b = jnp.zeros((16,))\nprint(make_jaxpr(layer)(x, W, b))\n1\n2\n3\n4\n5\n6\n7\n8\n9\njs\n{ lambda ; a:f32[4,8] b:f32[8,16] c:f32[16]. let\n    d:f32[4,16] = dot_general[\n      dimension_numbers=(([1], [0]), ([], []))\n      preferred_element_type=float32\n    ] a b\n    e:f32[1,16] = broadcast_in_dim[broadcast_dimensions=(1,)] c\n    f:f32[4,16] = add d e\n    g:f32[4,16] = max f 0.0:f32[]\n  in (g,) }\n```\n\nYour `@`\n\nbecame `dot_general`\n\n— the one general contraction primitive that every matmul, batched matmul, and tensor contraction in JAX lowers to. The `dimension_numbers=(([1],[0]),([],[]))`\n\nsays “contract axis 1 of the left operand with axis 0 of the right, no batch axes,” which is an ordinary matrix multiply. Your `+ b`\n\nbecame an explicit `broadcast_in_dim`\n\nfollowed by `add`\n\n, because the jaxpr does not hide broadcasting — it is a real operation with a real cost, so it gets a node.\n\nThis is the level JAX reasons at. Import `jax.lax`\n\n, `jax.nn`\n\n, `jax.random`\n\n, and `jax.scipy.special`\n\n, then sweep the heap for live `core.Primitive`\n\nobjects, and you find **285** of them. 3 That count includes plenty of internal machinery you will never type, but the order of magnitude is the point: a few hundred typed operations, and no Python left.\n\n## Tracing is abstract interpretation, not bytecode surgery\n\nHere is where JAX diverges from `torch.compile`\n\nat the mechanism level, and the contrast explains everything downstream. Dynamo intercepts CPython at the frame-evaluation layer and symbolically executes your bytecode, opcode by opcode; when it meets something it cannot model, it takes a graph break and falls back to the interpreter. 4 JAX does none of that. It never looks at your bytecode. It just calls your function — with fake arguments.\n\nWhen you trace `f`\n\n, JAX does not pass in an array. It passes in a `Tracer`\n\n: an object carrying an *abstract value*, and recording every primitive applied to it. Put a print inside a traced function and you can see exactly what flows through:\n\n```\n1\n2\n3\n4\n5\n6\npython\ndef probe(x):\n    print(type(x).__name__, \"|\", x.aval)\n    return x + 1\n\nmake_jaxpr(probe)(jnp.ones((2, 3)))\n# DynamicJaxprTracer | float32[2,3]\n```\n\nThe value moving through your function is a `DynamicJaxprTracer`\n\nwhose `aval`\n\nis `float32[2,3]`\n\n. Your Python executes once, top to bottom, during tracing, and every `jax.numpy`\n\ncall it makes appends a node to the jaxpr being built. That is why `x @ W`\n\nturns into `dot_general`\n\n— the tracer’s `__matmul__`\n\nrecords the primitive instead of computing anything.\n\nThe abstract value is the entire worldview of the compiler. It knows your array is `float32[2,3]`\n\n. It does not know, and never asked, what the numbers are. Hold onto that sentence; the last section of this post is about how much *else* the aval now knows.\n\n## The trace boundary: what JAX refuses to see\n\nNow the offending `if`\n\n. A tracer has a type but no concrete value, so the moment your code demands a concrete boolean from it, tracing cannot continue:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\npython\nfrom jax import jit\n\ndef bad(x):\n    if x > 0:          # needs a real bool; x is a tracer\n        return x\n    return -x\n\njit(bad)(jnp.float32(1.0))\n1\n2\n3\n4\n5\njax.errors.TracerBoolConversionError:\nAttempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function bad for jit. This concrete\nvalue was not available in Python because it depends on the value of the\nargument x.\n```\n\n`x > 0`\n\nproduces another tracer — a `bool[]`\n\n— and a Python `if`\n\nneeds to branch *now*, which means collapsing that tracer to `True`\n\nor `False`\n\n, which JAX cannot do without the data. So it stops. This is the exact spot where `torch.compile`\n\nwould graph-break, and the two systems make opposite choices. Dynamo chooses “keep running, compile what you can.” JAX chooses “the trace stays pure and complete, so make the control flow part of the IR or don’t compile at all.”\n\nI have come around to preferring JAX’s choice, with a caveat. A graph break is a silent performance cliff — your model still runs, just slower, and you find out months later reading a profile. `TracerBoolConversionError`\n\nis loud and immediate; it fails in your face on the first call. That honesty is worth a lot. The cost is real: you cannot sprinkle data-dependent Python control flow through your model and expect it to compile.\n\n## Control flow has to become a primitive\n\nThe fix is to lift the branch out of Python and into the IR:[5](#fn:5)\n\n```\n1\n2\n3\n4\n5\n6\n7\npython\nfrom jax import lax\n\ndef good(x):\n    # select the branch inside the IR instead of in Python\n    return lax.cond(x > 0, lambda v: v, lambda v: -v, x)\n\nprint(make_jaxpr(good)(jnp.float32(1.0)))\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\njs\n{ lambda ; a:f32[]. let\n    b:bool[] = gt a 0.0:f32[]\n    c:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b\n    d:f32[] = cond[\n      branches=(\n        { lambda ; e:f32[]. let f:f32[] = neg e in (f,) }\n        { lambda ; g:f32[]. let  in (g,) }\n      )\n    ] c a\n  in (d,) }\n```\n\nThe branch is now a `cond`\n\nprimitive whose two arms are themselves little jaxprs, and the predicate is a traced `bool[]`\n\nrather than a Python bool. Both arms are in the IR because JAX traced both of them, which you can watch directly:\n\n```\n1\n2\n3\n4\n5\n6\npython\nhits = []\ndef true_fn(v):  hits.append(\"true\");  return v * 10\ndef false_fn(v): hits.append(\"false\"); return v - 10\n\nmake_jaxpr(lambda x: lax.cond(x > 0, true_fn, false_fn, x))(jnp.float32(1.0))\nprint(hits)     # ['true', 'false']\n```\n\nBoth branch functions ran at trace time. That is the price of making control flow data-independent, and it has a consequence people trip over: if one arm of your `cond`\n\nallocates a 10 GB intermediate, that allocation is in the program whether or not the predicate ever selects it.\n\nTracing both arms is not the same as *running* both. An unbatched `cond`\n\nexecutes one branch. But batch the predicate — which is what happens the moment you `vmap`\n\na per-example decision — and the primitive disappears:\n\n```\n1\n2\nf = lambda x: lax.cond(x > 0, lambda v: v * 10, lambda v: v - 10, x)\nprint(make_jaxpr(vmap(f))(jnp.arange(4, dtype=jnp.float32)))\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n13\njs\n{ lambda ; a:f32[4]. let\n    b:bool[4] = gt a 0.0:f32[]\n    c:i32[4] = convert_element_type[new_dtype=int32 weak_type=False] b\n    d:bool[4] = eq c 0:i32[]\n    e:f32[4] = stop_gradient a\n    f:f32[4] = select_n d e a\n    g:f32[4] = sub f 10.0:f32[]\n    h:bool[4] = eq c 1:i32[]\n    i:f32[4] = stop_gradient a\n    j:f32[4] = select_n h i a\n    k:f32[4] = mul j 10.0:f32[]\n    l:f32[4] = select_n c g k\n  in (l,) }\n```\n\nNo `cond`\n\nanywhere. Both the `sub`\n\nand the `mul`\n\nrun on all four lanes and a `select_n`\n\npicks per lane, because there is no such thing as a divergent branch across a vector register. 5 Each lane pays for both arms. When one arm is an attention block and the other is\n\n`return x`\n\n, that is the whole cost of the expensive arm on every example, and it is invisible unless you read the jaxpr.Loops split the same way. `lax.while_loop`\n\nis for data-dependent iteration counts. `lax.scan`\n\nis for a fixed count, and it is the one to internalize, because it traces the body exactly once:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\npython\nn = []\ndef body(carry, x):\n    n.append(1)\n    return carry + x, carry\n\nprint(make_jaxpr(lambda c, xs: lax.scan(body, c, xs))(\n      jnp.float32(0.), jnp.arange(1000, dtype=jnp.float32)))\nprint(len(n), \"trace(s) of the body\")\n1\n2\n3\n4\n5\n6\n7\n8\n9\njs\n{ lambda ; a:f32[] b:f32[1000]. let\n    c:f32[] d:f32[1000] = scan[\n      jaxpr={ lambda ; e:f32[] f:f32[]. let g:f32[] = add e f in (g, e) }\n      length=1000\n      reverse=False\n      unroll=1\n    ] a b\n  in (c, d) }\n1 trace(s) of the body\n```\n\nOne equation, `length=1000`\n\n, one traced body. Write the same loop as a Python `for`\n\nand JAX will happily trace it — it just unrolls the whole thing into the jaxpr, and you pay for that in the compiler. Here is a 5,000-step recurrence over a 32×32 weight matrix, measured both ways on CPU:\n\n```\n1\n2\n3\ntop-level eqns    trace     compile\npython for            10000       0.11 s    36.92 s\nlax.scan                  1       0.00 s     0.02 s\n```\n\nThirty-seven seconds versus twenty milliseconds, for a program that computes the same thing. This is the single most common self-inflicted wound in JAX code, and it does not announce itself — nothing errors, your job just sits in XLA for a minute before every run. The rule to internalize: anything that decides the *shape or structure* of the computation from a value must become a primitive, because the trace only knows types, never values.\n\n## Purity is not optional\n\nThe trace also assumes your function is pure — output determined entirely by its inputs, no side effects, no hidden state. This is not a style preference; it is load-bearing. Your Python runs while JAX traces it, so any side effect fires only during tracing, once per trace, and not on the cached calls that follow. A `print`\n\ninside a jitted function prints while JAX traces it, then goes silent on every cached call after — baffling if you did not expect it. Mutating a global from inside a traced function is worse: the write happens at trace time against whatever the state was then, and the compiled program has no memory of it. (The trace-counting trick further down works precisely because of this.)\n\nThe most visible consequence is randomness, and it is the single most complained-about thing in JAX. NumPy’s `np.random.randn()`\n\nreads and advances a global RNG state — a side effect, invisible to the tracer, impossible to reproduce or parallelize deterministically. JAX cannot trace that, so it refuses to have it. Randomness is instead an explicit, pure function of a key you pass in:[6](#fn:6)\n\n```\n1\n2\n3\nkey = jax.random.key(0)\njax.random.normal(key, (3,))   # [ 1.6226422  2.0252647 -0.4335944 ]\njax.random.normal(key, (3,))   # identical — same key, same draw\n```\n\nSame key, same numbers, every time — a draw is a deterministic function of its key, which is what a pure trace demands. To get fresh randomness you advance the state yourself, by splitting a key into new independent keys:\n\n```\n1\n2\nk1, k2 = jax.random.split(key)\njax.random.normal(k1, (3,))    # [ 1.0040143 -0.9063372 -0.7481722 ]\n```\n\nPeople hate this until they need it, and then they stop. Because the randomness is threaded explicitly, a JAX computation is bit-for-bit reproducible no matter how it is parallelized — the key, not a hidden global, determines the result. The generator underneath is Threefry, a counter-based cipher-like construction designed so that the *n*-th draw is computable directly from the counter without stepping through the first *n−1*, which is exactly what makes a split-and-scatter design work across 512 chips.[7](#fn:7)\n\nAnd a key is a first-class typed value, not a special case bolted onto the API. Trace a draw and look at the binder:\n\n```\n1\nprint(make_jaxpr(lambda k: jax.random.normal(k, (3,)))(key))\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\njs\n{ lambda ; o:key<fry>[]. let\n    p:f32[3] = jit[\n      name=_normal\n      jaxpr={ lambda ; o:key<fry>[]. let\n          ...\n          f:u32[3] = random_bits[bit_width=32 shape=(3,)] a\n          g:u32[3] = shift_right_logical f 9:u32[]\n          h:u32[3] = or g 1065353216:u32[]\n          i:f32[3] = bitcast_convert_type[new_dtype=float32] h\n          ...\n```\n\nThe input type is `key<fry>[]`\n\n— an extended dtype naming the Threefry implementation — and the body is the bit-twiddling you would write by hand: draw 32 random bits, shift right by 9 to keep the 23 mantissa bits, `or`\n\nin the exponent for 1.0 (`0x3F800000`\n\n= 1065353216), bitcast to float, subtract one. A uniform in [0,1) built from integer ops, fully visible in the IR, with no hidden state anywhere. The whole design is a direct consequence of the tracing model: no global state means no global RNG.\n\n## Every transformation is a jaxpr rewrite\n\nThis is the part that makes JAX feel less like a library and more like a small compiler you drive from Python. Because tracing produces one uniform IR, every transformation is defined as a rewrite of that IR — and rewrites compose.\n\n** grad is a source transformation.** Reverse-mode autodiff traces your function, then produces a\n\n*new*jaxpr that computes the gradient.\n\nYou can read it:\n\n[8](#fn:8)\n\n```\n1\n2\n3\n4\n5\n6\npython\nfrom jax import grad\n\ndef g(x):\n    return jnp.sum(jnp.sin(x) ** 2)\n\nprint(make_jaxpr(grad(g))(jnp.ones((3,))))\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\njs\n{ lambda ; a:f32[3]. let\n    b:f32[3] = sin a\n    c:f32[3] = cos a\n    d:f32[3] = integer_pow[y=2] b\n    e:f32[3] = mul 2.0:f32[] b\n    _:f32[] = reduce_sum[axes=(0,) out_sharding=None] d\n    f:f32[3] = broadcast_in_dim 1.0:f32[]\n    g:f32[3] = mul f e\n    h:f32[3] = mul g c\n  in (h,) }\n```\n\nThe whole chain rule is sitting in the `let`\n\nblock. `b = sin a`\n\nand `c = cos a`\n\nare the residuals the backward pass needs; `e = 2·sin(a)`\n\n; `f`\n\nis the incoming cotangent of the sum, which is `1`\n\n; and the result is `h = f · e · c = 2·sin(a)·cos(a)`\n\n— the derivative of `sin²(x)`\n\n. Two details. The forward result `d = sin²(a)`\n\nand its `reduce_sum`\n\nare computed and then *dropped* — bound to `_`\n\n— because the gradient does not need them; they are dead code XLA will delete later. And the gradient jaxpr is just another jaxpr, which is why `grad(grad(f))`\n\nand `grad(jit(f))`\n\nwork without special machinery. For how the forward trace gets linearized and transposed into this, I wrote up the mechanics in [the autodiff post](/posts/autodiff-representations/).\n\n** vmap rewrites primitives, not Python.** Vectorization is not a loop over your function. It rewrites each primitive to a batched version by threading a batch axis through the jaxpr:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\n9\nphp\ndef dot(a, b):\n    return jnp.dot(a, b)                       # vectors -> scalar\n\nmake_jaxpr(dot)(jnp.ones(5), jnp.ones(5))\n#   dot_general[dimension_numbers=(([0], [0]), ([], []))]\n\nfrom jax import vmap\nmake_jaxpr(vmap(dot))(jnp.ones((32, 5)), jnp.ones((32, 5)))\n#   dot_general[dimension_numbers=(([1], [1]), ([0], [0]))]\n```\n\nLook at what changed. The unbatched dot contracts axis 0 with axis 0 and produces a scalar. Under `vmap`\n\n, `dimension_numbers`\n\ngained a batch entry — `([0],[0])`\n\nin the third slot — so axis 0 of each operand is now a batch dimension and the contraction moved to axis 1. There is no Python loop and no per-example overhead; `vmap`\n\nproduced a single batched `dot_general`\n\nthat XLA lowers to one kernel. This is why the advice in JAX is always “write the function for one example and `vmap`\n\nit,” where the equivalent advice in eager frameworks is “manually batch everything and pray the shapes line up.”\n\nThe rewrites are table-driven, and the tables are not the same size:\n\n| Registry | Entries |\n|---|---|\nMLIR lowering rules (`mlir._lowerings` ) | 251 |\nbatching rules (`batching.fancy_primitive_batchers` ) | 225 |\nJVP rules (`ad.primitive_jvps` ) | 200 |\ntranspose rules (`ad.primitive_transposes` ) | 84 |\n\nA primitive needs a lowering rule to run at all, a batching rule to survive `vmap`\n\n, a JVP rule to be differentiated forward, and a transpose rule for the linearized computation to be reversed. Those are four independent obligations, and nobody has discharged all four for every primitive. Reverse-mode `grad`\n\nis the narrowest gate: 84 transpose rules. When you hit a primitive with no rule, JAX does not silently degrade — it raises, names the primitive, and tells you which rule is missing. Same philosophy as the `if`\n\n.\n\nBecause all three transformations are jaxpr-to-jaxpr rewrites over the same IR, they stack in any order. `jit(grad(vmap(f)))`\n\nis not a special case; it is three rewrites composed.\n\n## Retracing, and the cache key that decides your compile time\n\n`jit`\n\nis the transformation that stages the jaxpr out to XLA and caches the compiled result. The cache key is the thing to memorize, because it is where JAX’s compile-time behavior comes from. JAX keys on the *abstract value* of each argument and nothing else:[9](#fn:9)\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\npython\ncount = {\"n\": 0}\n\n@jit\ndef f(x):\n    count[\"n\"] += 1          # side effect: runs only while tracing\n    return x * 2 + 1\n\nf(jnp.ones((4,), jnp.float32));  # count -> 1   (trace + compile)\nf(jnp.ones((4,), jnp.float32));  # count -> 1   (cache hit)\nf(jnp.ones((8,), jnp.float32));  # count -> 2   (new shape: retrace)\nf(jnp.ones((8,), jnp.int32));    # count -> 3   (new dtype: retrace)\n```\n\nSame aval, cache hit, no retrace. Change the shape or the dtype and JAX traces and compiles a fresh program. This is the same failure mode as recompilation thrashing in [the torch.compile post](/posts/how-torch-compile-works/), but the cause is cleaner and more predictable: there is no guard tree to reason about, just the aval. If your batch size varies every step, you recompile every step. The fixes are shape padding, bucketing to a handful of fixed sizes, or symbolic shape polymorphism via `jax.export`\n\n, which lowers one shape-generic program at the cost of giving XLA less to work with.[10](#fn:10)\n\nThe flip side is `static_argnums`\n\n, which tells `jit`\n\nto treat an argument as a compile-time Python constant rather than a traced value. That is mandatory when an argument determines a *shape*, since shapes must be concrete during tracing:\n\n```\n1\n2\n3\n4\n5\n6\n7\npython\nfrom functools import partial\n\n@partial(jit, static_argnums=0)\ndef g(n, x):\n    return jnp.broadcast_to(x, (n,))   # n must be concrete to be a shape\n\nprint(make_jaxpr(g, static_argnums=0)(3, jnp.float32(7.0)))\n1\n2\n3\n4\n5\n6\njs\n{ lambda ; a:f32[]. let\n    b:f32[3] = jit[\n      name=g\n      jaxpr={ lambda ; a:f32[]. let b:f32[3] = broadcast_in_dim a in (b,) }\n    ] a\n  in (b,) }\n```\n\nThe `3`\n\nis gone from the argument list and baked into the type of `b`\n\n. A different `n`\n\nis a different compiled program — static arguments are part of the cache key too. Use them for things that genuinely define structure (a number of layers, a boolean flag) and never for values that vary a lot, or you are back to recompiling constantly.\n\n## PyTrees, briefly\n\nOne more front-end mechanism, because it is everywhere. JAX transformations operate on *PyTrees* — arbitrarily nested tuples, lists, and dicts of arrays. 11 A model’s parameters are usually one big nested dict, and JAX treats it as a container it can flatten into a flat list of leaves plus a\n\n`treedef`\n\ndescribing the structure:\n\n```\n1\n2\n3\n4\npython\nimport jax.tree_util as jtu\nparams = {\"w\": jnp.ones((2, 2)), \"layers\": [jnp.zeros((3,)), jnp.ones((3,))]}\nleaves, treedef = jtu.tree_flatten(params)\n# 3 leaves; treedef = PyTreeDef({'layers': [*, *], 'w': *})\n```\n\nThis is why `grad(loss)(params)`\n\nreturns a gradient with the exact structure of `params`\n\n, and why `jit`\n\naccepts and returns nested dicts transparently. Flatten to leaves, transform the leaves, unflatten with the saved `treedef`\n\n. Unglamorous plumbing doing a lot of quiet work.\n\n## The aval grew places\n\nGo back to that binder: `a:f32[4,8]`\n\n. Dtype and shape, and I said the compiler knows nothing else. That was true for most of JAX’s life, and it stopped being true recently. Look at what `ShapedArray`\n\nactually declares:[12](#fn:12)\n\n```\n1\n2\n3\nclass ShapedArray(AbstractValue):\n  __slots__ = ['shape', 'dtype', 'weak_type', 'sharding', 'manual_axis_type',\n               'memory_space', '_stripped_weak_type', '__weakref__']\n```\n\n`sharding`\n\n. `memory_space`\n\n. The abstract value — the thing tracing propagates, the thing that keys the compilation cache — now carries where the array lives.\n\nTo see it, build a mesh with **explicit** axes. JAX classifies each mesh axis as `Auto`\n\n, `Explicit`\n\n, or `Manual`\n\n; under `Auto`\n\n, sharding is a hint the compiler propagates on its own and the type stays quiet, which is what you get by default and why you may never have noticed any of this. 13 Mark the axes\n\n`Explicit`\n\nand placement moves into the type:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\npython\nimport jax\nfrom jax.sharding import PartitionSpec as P, AxisType\n\njax.config.update('jax_num_cpu_devices', 8)\nmesh = jax.make_mesh((4, 2), ('x', 'y'),\n                     axis_types=(AxisType.Explicit, AxisType.Explicit))\n\nwith jax.sharding.set_mesh(mesh):\n    a = jax.device_put(jnp.zeros((8, 16)), jax.NamedSharding(mesh, P('x', 'y')))\n    b = jax.device_put(jnp.zeros((16, 4)), jax.NamedSharding(mesh, P('y', None)))\n    print(jax.typeof(a), jax.typeof(b))\n1\nfloat32[8@x,16@y] float32[16@y,4]\n```\n\n`f32[8@x,16@y]`\n\n: eight rows sharded over mesh axis `x`\n\n, sixteen columns over `y`\n\n. The `@`\n\nis not decoration in a debug print — it comes from the aval’s own formatter, which walks the shape and the `PartitionSpec`\n\ntogether and emits one string. 14 Placement is\n\n*in the type*, per axis.\n\nAnd a type is only interesting if something checks it. Multiply those two arrays:\n\n```\n1\nmake_jaxpr(lambda p, q: p @ q)(a, b)\n1\n2\n3\n4\njax._src.core.ShardingTypeError: Contracting dimensions are sharded and it\nis ambiguous how the output should be sharded. Please specify the output\nsharding via the `out_sharding` parameter.\nGot lhs_contracting_spec=('y',) and rhs_contracting_spec=('y',)\n```\n\nThat is a type error about *placement*, raised during tracing, before any device does anything. The shapes are fine — `(8,16) @ (16,4)`\n\ncontracts cleanly. What does not check is where the operands live: contracting a `y`\n\n-sharded axis against a `y`\n\n-sharded axis leaves the result needing a reduction whose placement JAX will not guess for you. Elementwise ops are checked the same way:\n\n```\n1\n2\nShardingTypeError: add got incompatible shardings for broadcasting:\n('x', 'y'), ('y', 'x').\n```\n\n`ShardingTypeError`\n\nis a real exception class in `jax/_src/core.py`\n\n, thirteen lines above `MemorySpace`\n\nand three hundred above `ShapedArray`\n\n, and it is raised from **53 sites** across the codebase. 15 This is a type checker for data placement, and it has the density of raise-sites of one.\n\nAnswer the question it asked and the trace goes through, with the placement of the result written into the binder:\n\n```\n1\n2\n3\n4\n5\npython\ndef f(p, q):\n    return jax.lax.dot_general(p, q, (((1,), (0,)), ((), ())),\n                               out_sharding=P('x', None))\n\nprint(make_jaxpr(f)(a, b))\n1\n2\n3\njs\n{ lambda ; a:f32[8@x,16@y] b:f32[16@y,4]. let\n    c:f32[8@x,4] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a b\n  in (c,) }\n```\n\nMemory space rides along in the same slot. `MemorySpace`\n\nis a three-case enum — `Device`\n\n, `Host`\n\n, `Any`\n\n— and the formatter prints anything that is not `Device`\n\nas an angle-bracketed prefix on the dtype:\n\n```\n1\n2\n3\nh = jax.device_put(jnp.zeros((8, 16)),\n                   jax.NamedSharding(mesh, P('x', 'y'), memory_kind='pinned_host'))\nprint(jax.typeof(h))\n1\nfloat32<host>[8@x,16@y]\n```\n\nRead `float32<host>[8@x,16@y]`\n\nas one type expression and it tells you the element type, the memory space, the extent of each axis, and the mesh axis each is distributed over. That is a long way from `f32[4,8]`\n\n, and the direction of travel is unambiguous: the things a compiler needs to know about *where* data lives are migrating out of side-channel annotations and into the type of the value.\n\nI think that direction is correct, and I do not think it is finished. `sharding`\n\nand `memory_space`\n\ngot in because sharding bugs were expensive enough — silently wrong numbers, or a 100× slowdown from an unplanned all-to-all — that “annotate and hope the partitioner agrees” stopped being tolerable. That argument is not special to sharding. It applies to every property of a value that the hardware treats as load-bearing and the type currently ignores: layout, alignment, which of a chip’s several memories a buffer sits in, whether a pointer is valid on the device you are about to dereference it from. Those are still, today, comments and conventions and runtime aborts. There is no principled line separating “in the type” from “beside the type” here — there is a line someone drew, and it has already moved twice.\n\nWhat that means for the [tour of XLA](/posts/a-tour-of-xla-where-mlir-lives/) and its [rigid price](/posts/xla-review-and-critique/) is that the front end is handing the compiler more than it used to. And what it means for the next post is that sharding deserves its own treatment: how a `PartitionSpec`\n\nbecomes a mesh assignment, where GSPMD ends and the newer Shardy partitioner begins, and what happens on the axes you leave as `Auto`\n\n— which is to say, the ones where the type stays quiet and the compiler decides for you.\n\n## References\n\n*Disclaimer: Researched and drafted with AI assistance (Gemini 3.1 Pro, Claude Opus 4.8 and Opus 5). Direction, technical judgment, and final edits are mine; every claim is traceable to the sources cited above.*\n\n**Reproduction environment.** All jaxprs, error messages, timings, and registry counts in this post are from JAX 0.11.0 on the CPU backend, Python 3.14, with`jax.config.update('jax_num_cpu_devices', 8)`\n\nfor the mesh examples. Source line references are to the same release. ([Link](https://pypi.org/project/jax/0.11.0/))[↩︎](#fnref:1)**JAX: Understanding jaxprs.** The jaxpr grammar, binders,`let`\n\n-bound equations, and how tracing with abstract values produces one. ([Link](https://docs.jax.dev/en/latest/jaxpr.html))[↩︎](#fnref:2)**Primitive count method.** After`import jax.lax, jax.nn, jax.random, jax.scipy.special`\n\n, sweeping`gc.get_objects()`\n\nfor instances of`jax._src.core.Primitive`\n\nyields 285 live objects. This includes primitives that exist only for internal staging (`addupdate`\n\n,`accum_grad_in_ref`\n\n, Pallas-side ops) and is an upper bound on the user-facing surface, not a count of documented operations.[↩︎](#fnref:3)**torch.compile / TorchDynamo — frame-evaluation capture.** The bytecode-interception design and the graph-break fallback that JAX declines to have. ([Link](https://pytorch.org/docs/stable/torch.compiler_dynamo_overview.html))[↩︎](#fnref:4)**JAX: Control flow and logical operators**, plus the`lax.cond`\n\ndocstring. The docs cover why data-dependent Python control flow must become`lax.cond`\n\n,`lax.scan`\n\n, or`lax.while_loop`\n\n. The batching behaviour is stated in the source: “In contrast with`jax.lax.select`\n\n, using`cond`\n\nindicates that only one of the two branches is executed (up to compiler rewrites and optimizations). However, when transformed with`jax.vmap`\n\nto operate over a batch of predicates,`cond`\n\nis converted to`jax.lax.select`\n\n. Both branches will be traced in all cases.” (`jax/_src/lax/control_flow/conditionals.py:212`\n\n, JAX 0.11.0.) ([Link](https://docs.jax.dev/en/latest/control-flow.html))[↩︎](#fnref:5)[↩︎](#fnref:5:1)2**JAX: Pseudorandom numbers.** Why global RNG state is incompatible with the traced, functional model, and the explicit key/`split`\n\ndesign that replaces it. ([Link](https://docs.jax.dev/en/latest/random-numbers.html))[↩︎](#fnref:6)**Salmon, Moraes, Dror, Shaw — “Parallel Random Numbers: As Easy as 1, 2, 3” (SC ‘11).** The counter-based Threefry generator underneath JAX’s`key<fry>`\n\ndtype, designed so the*n*-th output is computable directly from the counter. ([Link](https://dl.acm.org/doi/10.1145/2063384.2063405))[↩︎](#fnref:7)**JAX: Automatic differentiation.** How JAX composes forward-mode linearization and transposition to produce reverse-mode gradients as a program transformation over jaxprs. ([Link](https://docs.jax.dev/en/latest/automatic-differentiation.html))[↩︎](#fnref:8)**JAX: Just-in-time compilation.** Tracing, caching keyed on abstract shape and dtype, and`static_argnums`\n\n. ([Link](https://docs.jax.dev/en/latest/jit-compilation.html))[↩︎](#fnref:9)**JAX: Exporting and serialization — shape polymorphism.** Lowering a single shape-generic program with symbolic dimensions, and the constraints that come with it. ([Link](https://docs.jax.dev/en/latest/export/shape_poly.html))[↩︎](#fnref:10)**JAX: Working with pytrees.** Flatten/unflatten,`treedef`\n\n, and how transformations map over leaves. ([Link](https://docs.jax.dev/en/latest/working-with-pytrees.html))[↩︎](#fnref:11)`ShapedArray.__slots__`\n\n.`jax/_src/core.py:2450`\n\nin JAX 0.11.0 —`shape`\n\n,`dtype`\n\n,`weak_type`\n\n,`sharding`\n\n,`manual_axis_type`\n\n,`memory_space`\n\n. ([Link](https://github.com/jax-ml/jax/blob/main/jax/_src/core.py))[↩︎](#fnref:12)**JAX: Explicit sharding (sharding in types).** The`Auto`\n\n/`Explicit`\n\n/`Manual`\n\naxis-type distinction, defined at`jax/_src/mesh.py:111`\n\n, and what changes when an axis is marked`Explicit`\n\n. ([Link](https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html))[↩︎](#fnref:13)**The aval formatter.**`str_short_aval`\n\nat`jax/_src/core.py:2632`\n\ncomposes the dtype, an angle-bracketed memory space for anything other than`MemorySpace.Device`\n\n, and the shape interleaved with the`PartitionSpec`\n\nby`_get_shape_sharding_str`\n\nat`jax/_src/core.py:2611`\n\n. ([Link](https://github.com/jax-ml/jax/blob/main/jax/_src/core.py))[↩︎](#fnref:14)Declared at`ShardingTypeError`\n\n.`jax/_src/core.py:2197`\n\n;`grep -rn \"ShardingTypeError(\" jax/`\n\nover the installed 0.11.0 package returns 53 raise sites.`MemorySpace`\n\nis declared four lines below it at`jax/_src/core.py:2201`\n\n. ([Link](https://github.com/jax-ml/jax/blob/main/jax/_src/core.py))[↩︎](#fnref:15)\n\n[CC BY 4.0](https://creativecommons.org/licenses/by/4.0/)by the author.", "url": "https://wpnews.pro/news/what-jax-traces-and-what-it-refuses", "canonical_source": "https://hiraditya.github.io/posts/jax-tracing-and-jaxpr/", "published_at": "2026-07-28 19:00:00+00:00", "updated_at": "2026-07-29 07:32:07.033712+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools"], "entities": ["JAX", "PyTorch", "jax.jit", "torch.compile", "jax.make_jaxpr", "jax.lax", "jax.numpy", "jax.random"], "alternates": {"html": "https://wpnews.pro/news/what-jax-traces-and-what-it-refuses", "markdown": "https://wpnews.pro/news/what-jax-traces-and-what-it-refuses.md", "text": "https://wpnews.pro/news/what-jax-traces-and-what-it-refuses.txt", "jsonld": "https://wpnews.pro/news/what-jax-traces-and-what-it-refuses.jsonld"}}