We’re excited to share that TorchTPU—our native PyTorch stack for Google Cloud TPUs—is now open source on GitHub, alongside the early preview of the TPU Agent Suite!

Co-developed with the PyTorch team at Meta, TorchTPU brings Cloud TPUs (v5e, v5p, v6e, and v7 / Ironwood) alongside GPUs and CPUs as a native backend option for PyTorch. You can now train, fine-tune, and serve PyTorch models on TPUs with minimal code changes using the standard APIs and open-source tools you already rely on.

⚡ Quickstart: Installing TorchTPU

You can install the latest torch_tpu release directly via pip (Python 3.11+):

python -m venv .venv

source .venv/bin/activate

pip install torch_tpu --extra-index-url https://download.pytorch.org/whl/cpu

🛠️ What’s Included in TorchTPU?

⚬ Native PyTorch Experience: Debug models effortlessly in eager mode (.to("tpu")) and transition to production-grade performance with torch.compile(), backed by custom ATen kernels and an XLA compilation cache.

⚬ Out-of-the-Box Ecosystem…