7 / 8 JAX & XLA Mental Model

From Python to Silicon

When you call a jax.jit-decorated function for the first time with a given input shape and dtype, what does JAX do before it ever runs on the accelerator?

A

Serializes the function with pickle and sends it to a remote accelerator scheduler for compilation

B

Traces the Python function into a jaxpr, lowers it to HLO, and hands that to XLA to compile into a device-specific executable

C

Directly translates Python bytecode into CUDA kernels without any intermediate representation

D

Interprets the function eagerly once on CPU to check for errors, then copies the result to the device