XLA compilation

XLA (Accelerated Linear Algebra) takes an entire ML computation, described as a graph of array operations, and compiles it ahead of time into one optimized program for a specific accelerator — much as a C compiler turns source into a tuned binary, rather than interpreting operations one at a time. The payoff is whole-program optimization: it sees the graph globally and can rearrange it.

XLA lowers a high-level graph to its HLO intermediate representation, then runs target-independent passes — algebraic simplification, common-subexpression elimination, layout assignment, buffer allocation, and aggressive operator fusion — followed by backend codegen for GPU, TPU, or CPU. Because it knows static shapes ahead of time, it can pick layouts and tile sizes precisely and pack buffers tightly; the cost is that changing shapes triggers recompilation, so dynamic shapes and heavily data-dependent control flow are awkward. It is the compiler behind JAX and TensorFlow and the native programming model for TPUs.

For TPU-scale training and JAX workloads, XLA's fusion and layout choices are what close the gap to peak hardware utilization, and its sharding annotations (GSPMD) let one compiled program span thousands of chips. The tradeoff — recompilation and a more static world — is simply the price of ahead-of-time global optimization.