A developer’s guide to training with Ironwood TPUs

The transition toward trillion-parameter AI models has created an exponential demand for computational resources, testing the limits of traditional infrastructure. The seventh-generation Ironwood TPU features Google’s custom-designed AI infrastructure: It is engineered to scale as a holistic system supporting pods of up to 9,216 chips by combining Inter-Chip Interconnect (ICI), Optical Circuit Switch (OCS), Data Center Network (DCN) and massive aggregated High Bandwidth Memory (HBM) capacity. In addition, Ironwood features an integrated co-design between hardware architecture and software, introducing innovations such as Compiler-Centric XLA and Python-native kernels like Pallas and Mosaic. Together, these features significantly scale organizations’ capacity to train and serve sophisticated frontier models, optimize the entire AI lifecycle and enable sustained high performance. This technical overview explores the specific methods and tools within the JAX and MaxText ecosystems designe