Sponsored Session: TorchTPU: Expanding TPU Programmabil... Kat Ko, Claudio Basile & Jana van Greunen

This talk covers TorchTPU, a new initiative by Google aimed at expanding TPU programmability to PyTorch, which traditionally focused on TensorFlow and JAX. The speakers detail the evolution of TorchTPU, discussing the challenges of translating dynamic, eager execution into optimized TPU computations and the development of a native, eager-first PyTorch backend. They highlight key milestones such as native integration with torch.compile, DTensor, and support for the latest TPU architecture. These enhancements facilitate running multi-billion parameter models on TPUs with minimal code changes while allowing users to implement model-specific optimizations. The session concludes with insights into the roadmap for 2026.