Ray AI libraries run on TPU in Google Cloud

7/24/2026

Ray on TPU streamlines multi-host serving, JAX training, and data pipelines See how Serve, Data, and JaxTrainer simplify TPU workflows with Ray

Google Cloud has published part 2 of its guide to running Ray on TPU, focusing on the Ray AI libraries used to build workloads on top of the underlying slicebased placement model. The post explains how Ray Serve can deploy multihost models on TPUs with a single topology setting, how Ray Data's iterjaxbatches helps feed TPU jobs with devicesharded JAX arrays, and how JaxTrainer supports distributed JAX training with checkpointing and fault tolerance. It also notes new official Ray TPU Docker images, dashboard support for TPU metrics, and ongoing work to expand TPU integration across Ray tools.