Alongside TorchTPU, we’re releasing an early preview of the TPU agent suite on GitHub. Currently in preview for PyTorch inference workloads on Cloud TPUs, this open-source set of agents helps you onboard, profile, and serve models from inside the coding agent you already use. At launch, it works in Claude Code and Google Antigravity, with more coding agents to follow.
You talk to one front-end agent (tpu-agent), which routes each task to a specialist sub-agent: tpu-model-onboarding helps bring a model up, benchmark it, or serve it; tpu-diagnostics captures and inspects profiles to find what is slowing a model down; tpu-kernel-development helps write and tune custom Pallas kernels; and tpu-knowledge-base answers hardware and software questions with a source you can check. Before onboarding runs, it shows you a plan for approval, and the serving workflow ends with a numerical parity check on the TPU against the original PyTorch model on CPU.
The agents' actions are grounded, bounded, and checked, and you stay in the loop. The suite is open source and in preview on GitHub under google-pytorch.
Chapters:
0:00 - Introduction
0:11 - TPUs, designed for machine learning
0:24 - TorchTPU: native PyTorch on Cloud TPU
0:35 - The TPU agent suite
0:47 - Works inside your coding agent
1:17 - From running to running well
1:44 - One front end, the right specialist for each job
2:29 - A serving example: the same checklist every time
3:04 - Three design choices
3:14 - Grounded: answers come with a source
3:23 - Bounded: a budget, then it hands back
3:41 - Checked: parity with plain PyTorch before it is called done
3:59 - Recap and getting started
🔗 Resources:
* Explore the TPU agent suite on GitHub → https://g.dev/cloud/tpu-agent-suite
* Explore TorchTPU on GitHub → https://g.dev/cloud/torch-tpu
* PyTorch on Cloud TPU → https://g.dev/cloud/pytorch-on-tpu
* Learn more about Google Cloud TPU → https://g.dev/cloud/tensor-processing-units
Speaker: Rob Mulla
Products Mentioned: Tensor Processing Units
You talk to one front-end agent (tpu-agent), which routes each task to a specialist sub-agent: tpu-model-onboarding helps bring a model up, benchmark it, or serve it; tpu-diagnostics captures and inspects profiles to find what is slowing a model down; tpu-kernel-development helps write and tune custom Pallas kernels; and tpu-knowledge-base answers hardware and software questions with a source you can check. Before onboarding runs, it shows you a plan for approval, and the serving workflow ends with a numerical parity check on the TPU against the original PyTorch model on CPU.
The agents' actions are grounded, bounded, and checked, and you stay in the loop. The suite is open source and in preview on GitHub under google-pytorch.
Chapters:
0:00 - Introduction
0:11 - TPUs, designed for machine learning
0:24 - TorchTPU: native PyTorch on Cloud TPU
0:35 - The TPU agent suite
0:47 - Works inside your coding agent
1:17 - From running to running well
1:44 - One front end, the right specialist for each job
2:29 - A serving example: the same checklist every time
3:04 - Three design choices
3:14 - Grounded: answers come with a source
3:23 - Bounded: a budget, then it hands back
3:41 - Checked: parity with plain PyTorch before it is called done
3:59 - Recap and getting started
🔗 Resources:
* Explore the TPU agent suite on GitHub → https://g.dev/cloud/tpu-agent-suite
* Explore TorchTPU on GitHub → https://g.dev/cloud/torch-tpu
* PyTorch on Cloud TPU → https://g.dev/cloud/pytorch-on-tpu
* Learn more about Google Cloud TPU → https://g.dev/cloud/tensor-processing-units
Speaker: Rob Mulla
Products Mentioned: Tensor Processing Units