{ "description": "Browse aymeric's, an independent, searchable archive of readable snapshots from followed feeds, preserved in Git with aggr.", "feed_url": "https://aggr.aymericbeaumet.com/sources/github.com/google-pytorch/feed.json", "home_page_url": "https://aggr.aymericbeaumet.com/sources/github.com/google-pytorch/", "items": [ { "authors": [ { "name": "ilreb" } ], "content_html": "

\"TorchTPU

\n

Description

\n

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.

\n

Features

\n\n

User Guide and Interactive Notebooks

\n

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.

\n

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.

\n

Installation

\n

The same instructions are available in the user guide linked above.

\n

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.

\n

First create a Python virtual environment:

\n
python -m venv .venv\nsource .venv/bin/activate\n
\n

or, with uv

\n
uv venv .venv\nsource .venv/bin/activate\n
\n

Install via pip

\n

Install the latest public release (0.2.0 on PyPI) (for uv prepend uv to all pip commands):

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

Build from source

\n

Wheels can be built via:

\n
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel\n
\n

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:

\n
bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \\\n    --//:allow_unoptimized_wheel=True\n
\n

--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.

\n

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:

\n
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe\n
\n

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.

\n

Install wheel via:

\n
cd <path_to_repo>\npython -m venv .venv; source .venv/bin/activate\n# The extra-index-url provides the CPU version of torch rather than the default CUDA build.\npip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu\n
\n

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 shown above or install the CPU build separately.

\n

Getting Started

\n

Some pointers to get you started:

\n\n

Dependencies

\n

The project uses Bazel for dependency management. Key Python dependencies are:

\n

TBD

\n

Bazel setup

\n

This project uses Bazel for building and dependency management.

\n

To install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.

\n

There are several ways to install Bazelisk. Here are some common ways:

\n\n

See the Bazelisk README for more install instructions

\n

However you install it, add it's location to PATH.

\n

To verify the install, run bazel info in the repo's root directory:

\n
bazel info\n
\n

It should print information about the repo without errors.

\n

VSCode clangd Setup

\n

For C++ code navigation, it is recommended to use VSCode with clangd extension. Relying on a compilation database generated from Bazel action graph, clangd is able to enrich the editor with various smart features, including code completions, go-to definitions, etc.

\n
    \n
  1. \n

    Install VSCode extensions via UI or the following command:

    \n
    code --install-extension llvm-vs-code-extensions.vscode-clangd\n# Uninstall Microsoft C++ extensions to avoid interference.\ncode --uninstall-extension ms-vscode.cpptools\n
    \n
    2. \n
  2. \n

    Generate the initial compile_commands.json:

    \n
    tools/setup_clangd.py\n
    \n
    3. \n
\n

It may take a while (> 30 minutes) for the first run to make a full build, but reruns will be incremental and cheap. If you have run bazel build recently, generated files already exist so you can pass --no-build to skip waiting for heavy builds (note that on a fresh workspace, skipping the build may prevent clangd from resolving generated files). See more options in the script file-level comments.

\n

Alternatively, you can invoke the script via nox (arguments after -- are forwarded):

\n
nox -s refresh_compile_commands -- --no-build\nnox -s refresh_compile_commands -- //src/torch_tpu/ops/...\n
\n

Contributing

\n

We welcome contributions to the PyTorch on TPU project! Please see our CONTRIBUTING.md for guidelines on how to contribute, including signing our Contributor License Agreement (CLA) and code review processes.

\n

Star History

\n

\"Star

\n

License

\n

This project is licensed under the Apache 2.0 License. See the LICENSE file for details.

\n", "date_modified": "2026-10-20T16:40:24+00:00", "date_published": "2026-10-20T16:40:24+00:00", "external_url": "https://github.com/google-pytorch/torch_tpu", "id": "urn:aggr:item:ec897b5058ffe73a88e095b3f319f1a09c436003", "image": "https://opengraph.githubassets.com/cd86c84e93f7b104a4d7723deb52a1b58d9418986b422db17cdc92ee64b72b49/google-pytorch/torch_tpu", "summary": "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…", "tags": [ "news" ], "title": "Google open-sources TorchTPU, a native PyTorch back end for TPUs", "url": "https://aggr.aymericbeaumet.com/items/hnrss-org/2026-10-20-google-open-sources-torchtpu-a-native-pytorch-back-end-for/" } ], "language": "en", "title": "github.com/google-pytorch", "version": "https://jsonfeed.org/version/1.1" }