google-pytorch / torch_tpu

## Description
This project provides a PyTorch backend for Google's Tensor Processing Units
(TPUs), enabling users to run PyTorch models and operations on TPU devices. It
includes custom ATen kernels for various PyTorch operations, a compilation cache
for optimizing execution, and examples for distributed training and model
inference.
## Features
* PyTorch ATen kernel implementation for TPU.
* Compilation cache for optimized execution of XLA computations.
* Support for distributed training (Data Parallelism, Tensor Parallelism).
* Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).
* Integration with `torch.compile()` for graph mode execution.
* Utilities for debugging and benchmarking.
## User Guide and Interactive Notebooks
A user guide is available at
https://torchtpu.readthedocs.io/. This is a
static version of the interactive notebooks. They are available in the docs
folder.
The notebooks are Marimo (similar to Jupyter). They can be run with `marimo edit
file.py` and then you can access them in your browser. You will need to be
running this from a TPU VM, a v6e GCE instance is easiest.
## Installation
The same instructions are available in the user guide linked above.
TorchTPU's Python version support policy matches
PyTorch's Release Compatibility Matrix.
Currently, **Python 3.11+** is supported. See the
Python Version Support Policy for
more details.
First create a Python virtual environment:
 [code] 
or, with `uv`
 [code] 
### Install via pip
Install the latest public release
(0.2.0 on PyPI) (for `uv` prepend
`uv` to all `pip` commands):
 [code] 
### Build from source
Wheels can be built via:
 [code] 
All OSS Bazel builds default to optimized mode (`-c opt`) automatically via
`.bazelrc`. Unoptimized builds (`-c fastbuild`, `-c dbg`) can degrade
performance and are not used for wheel builds. If you are developing or
debugging C++ code and explicitly want an unoptimized wheel, pass
`--//:allow_unoptimized_wheel=True`:
 [code] 
`--config=wheel_common` is required, and the build fails at build time without
it. The config sets `--//:wheel_build=True`, which routes the shared XLA/MLIR
backend through the `torch_version`-reset transition so that pywrap factors it
into a single `libxla_base.so`; built without the flag, every per-version common
ships its own copy of the backend and the resulting wheel aborts on `import
torch` with duplicate static registrations.
`wheel_common` also turns on the Remote Build Execution (RBE) remote cache and
remote executor. If you do not have RBE credentials, append `--config=no_rbe` to
strip those flags back out:
 [code] 
The order matters here: `--config=no_rbe` has to come after
`--config=wheel_common`, because the later config is the one that wins for the
remote-cache and remote-executor flags they both set.
Install wheel via:
 [code] 
Note that installing without `--extra-index-url
https://download.pytorch.org/whl/cpu` will install the CUDA version of `torch`;
pass `--extra-index-url` as show