---
title: "Run Ray on TPU, Part 2: Ray AI libraries"
description: "Running Ray on TPUs requires multi-host workloads to remain on one intact TPU slice whose hosts share a high-speed ICI link. This guide presents Ray's AI libraries on top of GKE and Ray Core: Ray Serve for serving, Ray Data's iter_jax_batches() for input pipelines and batch inference, and Ray Train's JaxTrainer for distributed JAX training. Serve uses accelerator_config.topology to defer slice placement to the replica, while iter_jax_batches() returns device-sharded JAX arrays and handles ragged final batches through drop, pad, or raise choices. JaxTrainer takes a training function and slice topology, launches one worker per host, and provides checkpointing, fault-tolerant restarts, and multi-slice scale-out; official TPU images and dashboard metrics round out the workflow."
---

# Run Ray on TPU, Part 2: Ray AI libraries

[Google](https://yomu.fyi/company/google) · Ivan Nardini, Spencer Peterson · Jul 24, 2026

**Type:** Tutorial

## Summary

Running Ray on TPUs requires multi-host workloads to remain on one intact TPU slice whose hosts share a high-speed ICI link. This guide presents Ray's AI libraries on top of GKE and Ray Core: Ray Serve for serving, Ray Data's iter\_jax\_batches() for input pipelines and batch inference, and Ray Train's JaxTrainer for distributed JAX training. Serve uses accelerator\_config.topology to defer slice placement to the replica, while iter\_jax\_batches() returns device-sharded JAX arrays and handles ragged final batches through drop, pad, or raise choices. JaxTrainer takes a training function and slice topology, launches one worker per host, and provides checkpointing, fault-tolerant restarts, and multi-slice scale-out; official TPU images and dashboard metrics round out the workflow.

## Context

Multi-host TPU models must land on one whole slice because workers on different slices cannot communicate over the ICI; otherwise jobs can hang. TPU input pipelines can also become bottlenecks when data loading does not keep the accelerators fed, while training previously required accounting for slice topology in application code.

## Approach / What changed

Use GKE with the Ray Operator and Ray Core's slice\_placement\_group() to reserve a declared TPU topology, then configure Ray Serve with accelerator\_config.topology, Ray Data with iter\_jax\_batches(), or Ray Train with JaxTrainer, use\_tpu=True, and a topology. The guide also recommends RayService for production serving and describes official -tpu images and dashboard and profiler support.

## Takeaways

- Without accelerator\_config.topology, Serve can place per-chip bundles across separate slices, leaving a multi-host deployment stuck in DEPLOYING because the workers cannot complete their first collective.
- JAX should be imported inside train\_loop\_per\_worker because each worker initializes JAX in its own TPU context; JaxTrainer's topology value describes the slice shape, not a chip count.
- Official Ray TPU images include the JAX/TPU stack and profiling tools, while the Ray Dashboard shows TPU utilization and memory and can attach to per-worker JAX profiling.

**Tags:** [JAX](https://yomu.fyi/topic/jax), [Ray](https://yomu.fyi/topic/ray), [TPU](https://yomu.fyi/topic/tpu)

- Source: [Google](https://developers.googleblog.com/run-ray-on-tpu-part-2-ray-ai-libraries)
- Source URL: https://developers.googleblog.com/run-ray-on-tpu-part-2-ray-ai-libraries
- Ingested by Yomu: 2026-09-09T14:28:22.138Z

[Read original post](https://developers.googleblog.com/run-ray-on-tpu-part-2-ray-ai-libraries)