4 / 8 JAX & XLA Mental Model

Traced Once

jax.jit traces a Python function once per input signature (shapes/dtypes/pytree structure) and caches the resulting XLA computation. If the traced function contains a bare Python if x > 0: where x is a traced (non-static) array, what happens?

A

JAX lowers any array-valued Python if into an XLA select/cond op automatically

B

JAX silently follows the branch matching the sign bit of the abstract value, producing a result without erroring

C

JAX raises a TracerBoolConversionError/ConcretizationTypeError because a traced value has no concrete boolean at trace time

D

JAX evaluates the branch eagerly using the runtime value and recompiles automatically whenever that value flips sign

E

JAX raises a shape mismatch error because comparison ops must be pre-registered as primitives