JAX, XLA, and TPU Compilation

3.14.1JAX Tracing, jaxpr, and XLA#

3.14.2StableHLO, HLO, and Multi-Stage Optimization#

3.14.3Sharding, GSPMD/Shardy, and Collective Insertion#

3.14.4Layout Assignment, Buffer Assignment, and Memory Scheduling#

3.14.5Cooperation Between Pallas/Mosaic and High-Level Graph Compilation#

3.14.6MXU/Vector/DMA Cooperation and Static Scheduling#

3.14.7The libtpu Delivery Boundary and the Scope of Public TPU ISA Information#