{
  "@context": "https://schema.org",
  "@id": "https://aggr.aymericbeaumet.com/items/hnrss-org/2026-10-20-google-open-sources-torchtpu-a-native-pytorch-back-end-for/#webpage",
  "@type": [
    "WebPage",
    "ArchiveComponent"
  ],
  "articleBody": "[![TorchTPU Logo](https://github.com/google-pytorch/torch_tpu/raw/main/docs/images/torchtpu_logo.avif)](https://github.com/google-pytorch/torch_tpu/blob/main/docs/images/torchtpu_logo.avif)\n\n## Description\n\nThis 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\n## Features\n\n- PyTorch ATen kernel implementation for TPU.\n- Compilation cache for optimized execution of XLA computations.\n- Support for distributed training (Data Parallelism, Tensor Parallelism).\n- Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).\n- Integration with `torch.compile()` for graph mode execution.\n- Utilities for debugging and benchmarking.\n\n## User Guide and Interactive Notebooks\n\nA user guide is available at [https://torchtpu.readthedocs.io/](https://torchtpu.readthedocs.io/). This is a static version of the interactive notebooks. They are available in the docs folder.\n\nThe 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\n## Installation\n\nThe same instructions are available in the user guide linked above.\n\nTorchTPU's Python version support policy matches [PyTorch's Release Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix). Currently, **Python 3.11+** is supported. See the [Python Version Support Policy](https://github.com/google-pytorch/torch_tpu/blob/main/docs/developers/python_support_policy.md) for more details.\n\nFirst create a Python virtual environment:\n\n```shell\npython -m venv .venv\nsource .venv/bin/activate\n```\n\nor, with `uv`\n\n```shell\nuv venv .venv\nsource .venv/bin/activate\n```\n\n### Install via pip\n\nInstall the latest public release ([0.2.0 on PyPI](https://pypi.org/project/torch-tpu/0.2.0/)) (for `uv` prepend `uv` to all `pip` commands):\n\n```shell\npip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu\n```\n\n### Build from source\n\nWheels can be built via:\n\n```shell\nbazel build --config=wheel_common //ci/wheel:torch_tpu_wheel\n```\n\nAll 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\n```shell\nbazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \\\n    --//:allow_unoptimized_wheel=True\n```\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\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\n```shell\nbazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe\n```\n\nThe 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\nInstall wheel via:\n\n```shell\ncd <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\nNote 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\n## Getting Started\n\nSome pointers to get you started:\n\n- [PyTorch: Official PyTorch documentation](https://github.com/pytorch/pytorch)\n- [Tutorials and examples: easy to understand PyTorch code using TorchTPU](https://torchtpu.readthedocs.io/)\n- [Public APIs: the TorchTPU APIs that you can use](https://github.com/google-pytorch/torch_tpu/blob/main/docs/public_apis.md)\n\n## Dependencies\n\nThe project uses Bazel for dependency management. Key Python dependencies are:\n\nTBD\n\n## Bazel setup\n\nThis project uses Bazel for building and dependency management.\n\nTo install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.\n\nThere are several ways to install Bazelisk. Here are some common ways:\n\n- If you have Go installed:\n\n  ```\n  go install github.com/bazelbuild/bazelisk@latest\n  ```\n\n- Download the binary directly:\n\n  ```\n  wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64\n  ```\n\nSee the [Bazelisk README](https://github.com/bazelbuild/bazelisk/blob/master/README.md) for more install instructions\n\nHowever you install it, add it's location to `PATH`.\n\nTo verify the install, run `bazel info` in the repo's root directory:\n\n```\nbazel info\n```\n\nIt should print information about the repo without errors.\n\n## VSCode `clangd` Setup\n\nFor C++ code navigation, it is recommended to use VSCode with [`clangd` extension](https://marketplace.visualstudio.com/items?itemName=llvm-vs-code-extensions.vscode-clangd). Relying on a [compilation database](https://clang.llvm.org/docs/JSONCompilationDatabase.html) generated from Bazel action graph, [clangd](https://clangd.llvm.org/) is able to enrich the editor with various smart features, including code completions, go-to definitions, *etc.*\n\n1. Install VSCode extensions via UI or the following command:\n\n   ```shell\n   code --install-extension llvm-vs-code-extensions.vscode-clangd\n   # Uninstall Microsoft C++ extensions to avoid interference.\n   code --uninstall-extension ms-vscode.cpptools\n   ```\n\n2. Generate the initial `compile_commands.json`:\n\n   ```shell\n   tools/setup_clangd.py\n   ```\n\nIt 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](https://github.com/google-pytorch/torch_tpu/blob/main/tools/setup_clangd.py) file-level comments.\n\nAlternatively, you can invoke the script via `nox` (arguments after `--` are forwarded):\n\n```shell\nnox -s refresh_compile_commands -- --no-build\nnox -s refresh_compile_commands -- //src/torch_tpu/ops/...\n```\n\n## Contributing\n\nWe welcome contributions to the PyTorch on TPU project! Please see our [CONTRIBUTING.md](https://github.com/google-pytorch/torch_tpu/blob/main/CONTRIBUTING.md) for guidelines on how to contribute, including signing our Contributor License Agreement (CLA) and code review processes.\n\n## Star History\n\n[![Star History Chart](https://camo.githubusercontent.com/5dd8f8b127ebb9d4b8e81870ade5b9dea014587bc270e5ff7c4984ca8bf88d43/68747470733a2f2f6170692e737461722d686973746f72792e636f6d2f63686172743f7265706f733d676f6f676c652d7079746f7263682f746f7263685f74707526747970653d64617465266c6567656e643d746f702d6c656674)](https://www.star-history.com/?repos=google-pytorch%2Ftorch_tpu&type=date&legend=top-left)\n\n## License\n\nThis project is licensed under the Apache 2.0 License. See the [LICENSE](https://github.com/google-pytorch/torch_tpu/blob/main/LICENSE) file for details.\n",
  "authors": [
    "ilreb"
  ],
  "captured_at": "2026-10-20T21:31:53.761078291+00:00",
  "category": "news",
  "content_html": "<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/docs/images/torchtpu_logo.avif\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://github.com/google-pytorch/torch_tpu/raw/main/docs/images/torchtpu_logo.avif\" alt=\"TorchTPU Logo\" loading=\"eager\" decoding=\"async\" fetchpriority=\"high\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<h2>Description</h2>\n<p>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.</p>\n<h2>Features</h2>\n<ul>\n<li>PyTorch ATen kernel implementation for TPU.</li>\n<li>Compilation cache for optimized execution of XLA computations.</li>\n<li>Support for distributed training (Data Parallelism, Tensor Parallelism).</li>\n<li>Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).</li>\n<li>Integration with <code>torch.compile()</code> for graph mode execution.</li>\n<li>Utilities for debugging and benchmarking.</li>\n</ul>\n<h2>User Guide and Interactive Notebooks</h2>\n<p>A user guide is available at <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://torchtpu.readthedocs.io/\">https://torchtpu.readthedocs.io/</a>. This is a static version of the interactive notebooks. They are available in the docs folder.</p>\n<p>The notebooks are Marimo (similar to Jupyter). They can be run with <code>marimo edit file.py</code> 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.</p>\n<h2>Installation</h2>\n<p>The same instructions are available in the user guide linked above.</p>\n<p>TorchTPU's Python version support policy matches <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix\">PyTorch's Release Compatibility Matrix</a>. Currently, <strong>Python 3.11+</strong> is supported. See the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/docs/developers/python_support_policy.md\">Python Version Support Policy</a> for more details.</p>\n<p>First create a Python virtual environment:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">python</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>m</span> venv .venv</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-support syntax-function syntax-source syntax-shell\">source</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> .venv/bin/activate</span>\n</span></span></code></pre>\n<p>or, with <code>uv</code></p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">uv</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> venv .venv</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-support syntax-function syntax-source syntax-shell\">source</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> .venv/bin/activate</span>\n</span></span></code></pre>\n<h3>Install via pip</h3>\n<p>Install the latest public release (<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://pypi.org/project/torch-tpu/0.2.0/\">0.2.0 on PyPI</a>) (for <code>uv</code> prepend <code>uv</code> to all <code>pip</code> commands):</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">pip</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> install torch_tpu==0.2.0<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>extra-index-url</span> https://download.pytorch.org/whl/cpu</span>\n</span></span></code></pre>\n<h3>Build from source</h3>\n<p>Wheels can be built via:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">bazel</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> build<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-option syntax-shell\">=</span>wheel_common //ci/wheel:torch_tpu_wheel</span>\n</span></span></code></pre>\n<p>All OSS Bazel builds default to optimized mode (<code>-c opt</code>) automatically via <code>.bazelrc</code>. Unoptimized builds (<code>-c fastbuild</code>, <code>-c dbg</code>) 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 <code>--//:allow_unoptimized_wheel=True</code>:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">bazel</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> build<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>c</span> dbg<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-option syntax-shell\">=</span>wheel_common //ci/wheel:torch_tpu_wheel <span class=\"syntax-punctuation syntax-separator syntax-continuation syntax-line syntax-shell\">\\\n</span>    --//:allow_unoptimized_wheel=True</span>\n</span></span></code></pre>\n<p><code>--config=wheel_common</code> is required, and the build fails at build time without it. The config sets <code>--//:wheel_build=True</code>, which routes the shared XLA/MLIR backend through the <code>torch_version</code>-reset transition so that pywrap factors it into a single <code>libxla_base.so</code>; built without the flag, every per-version common ships its own copy of the backend and the resulting wheel aborts on <code>import torch</code> with duplicate static registrations.</p>\n<p><code>wheel_common</code> also turns on the Remote Build Execution (RBE) remote cache and remote executor. If you do not have RBE credentials, append <code>--config=no_rbe</code> to strip those flags back out:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">bazel</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> build<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-option syntax-shell\">=</span>wheel_common //ci/wheel:torch_tpu_wheel<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>config</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-option syntax-shell\">=</span>no_rbe</span>\n</span></span></code></pre>\n<p>The order matters here: <code>--config=no_rbe</code> has to come after <code>--config=wheel_common</code>, because the later config is the one that wins for the remote-cache and remote-executor flags they both set.</p>\n<p>Install wheel via:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-support syntax-function syntax-cd syntax-shell\">cd</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&lt;</span>path_to_repo<span class=\"syntax-keyword syntax-operator syntax-assignment syntax-redirection syntax-shell\">&gt;</span>\n</span><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">python</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>m</span> venv .venv</span><span class=\"syntax-keyword syntax-operator syntax-logical syntax-continue syntax-shell\">;</span> <span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-support syntax-function syntax-source syntax-shell\">source</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> .venv/bin/activate</span>\n<span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-begin syntax-shell\">#</span></span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\"> The extra-index-url provides the CPU version of torch rather than the default CUDA build.</span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\">\n</span><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">pip</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> install ../bazel-bin/ci/wheel/<span class=\"syntax-keyword syntax-operator syntax-regexp syntax-quantifier syntax-shell\">*</span>.whl<span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>extra-index-url</span> https://download.pytorch.org/whl/cpu</span>\n</span></span></code></pre>\n<p>Note that installing without <code>--extra-index-url https://download.pytorch.org/whl/cpu</code> will install the CUDA version of <code>torch</code>; pass <code>--extra-index-url</code> as shown above or install the CPU build separately.</p>\n<h2>Getting Started</h2>\n<p>Some pointers to get you started:</p>\n<ul>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/pytorch/pytorch\">PyTorch: Official PyTorch documentation</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://torchtpu.readthedocs.io/\">Tutorials and examples: easy to understand PyTorch code using TorchTPU</a></li>\n<li><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/docs/public_apis.md\">Public APIs: the TorchTPU APIs that you can use</a></li>\n</ul>\n<h2>Dependencies</h2>\n<p>The project uses Bazel for dependency management. Key Python dependencies are:</p>\n<p>TBD</p>\n<h2>Bazel setup</h2>\n<p>This project uses Bazel for building and dependency management.</p>\n<p>To install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.</p>\n<p>There are several ways to install Bazelisk. Here are some common ways:</p>\n<ul>\n<li>\n<p>If you have Go installed:</p>\n<pre><code><span class=\"code-snippet\">go install github.com/bazelbuild/bazelisk@latest\n</span></code></pre>\n</li>\n<li>\n<p>Download the binary directly:</p>\n<pre><code><span class=\"code-snippet\">wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64\n</span></code></pre>\n</li>\n</ul>\n<p>See the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/bazelbuild/bazelisk/blob/master/README.md\">Bazelisk README</a> for more install instructions</p>\n<p>However you install it, add it's location to <code>PATH</code>.</p>\n<p>To verify the install, run <code>bazel info</code> in the repo's root directory:</p>\n<pre><code><span class=\"code-snippet\">bazel info\n</span></code></pre>\n<p>It should print information about the repo without errors.</p>\n<h2>VSCode <code>clangd</code> Setup</h2>\n<p>For C++ code navigation, it is recommended to use VSCode with <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://marketplace.visualstudio.com/items?itemName=llvm-vs-code-extensions.vscode-clangd\"><code>clangd</code> extension</a>. Relying on a <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://clang.llvm.org/docs/JSONCompilationDatabase.html\">compilation database</a> generated from Bazel action graph, <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://clangd.llvm.org/\">clangd</a> is able to enrich the editor with various smart features, including code completions, go-to definitions, <em>etc.</em></p>\n<ol>\n<li>\n<p>Install VSCode extensions via UI or the following command:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">code</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>install-extension</span> llvm-vs-code-extensions.vscode-clangd</span>\n<span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-comment syntax-begin syntax-shell\">#</span></span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\"> Uninstall Microsoft C++ extensions to avoid interference.</span><span class=\"syntax-comment syntax-line syntax-number-sign syntax-shell\">\n</span><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">code</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> --</span>uninstall-extension</span> ms-vscode.cpptools</span>\n</span></span></code></pre>\n</li>\n<li>\n<p>Generate the initial <code>compile_commands.json</code>:</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">tools/setup_clangd.py</span></span>\n</span></span></code></pre>\n</li>\n</ol>\n<p>It may take a while (&gt; 30 minutes) for the first run to make a full build, but reruns will be incremental and cheap. If you have run <code>bazel build</code> recently, generated files already exist so you can pass <code>--no-build</code> to skip waiting for heavy builds (note that on a fresh workspace, skipping the build may prevent <code>clangd</code> from resolving generated files). See more options in the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/tools/setup_clangd.py\">script</a> file-level comments.</p>\n<p>Alternatively, you can invoke the script via <code>nox</code> (arguments after <code>--</code> are forwarded):</p>\n<pre><code class=\"language-shell\"><span class=\"code-snippet\" data-language=\"Shell\"><span class=\"syntax-source syntax-shell syntax-bash\"><span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">nox</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>s</span> refresh_compile_commands<span class=\"syntax-keyword syntax-operator syntax-end-of-options syntax-shell\"> --</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> --no-build</span>\n<span class=\"syntax-meta syntax-function-call syntax-shell\"><span class=\"syntax-variable syntax-function syntax-shell\">nox</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"><span class=\"syntax-variable syntax-parameter syntax-option syntax-shell\"><span class=\"syntax-punctuation syntax-definition syntax-parameter syntax-shell\"> -</span>s</span> refresh_compile_commands<span class=\"syntax-keyword syntax-operator syntax-end-of-options syntax-shell\"> --</span></span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-shell\"> //src/torch_tpu/ops/...</span>\n</span></span></code></pre>\n<h2>Contributing</h2>\n<p>We welcome contributions to the PyTorch on TPU project! Please see our <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/CONTRIBUTING.md\">CONTRIBUTING.md</a> for guidelines on how to contribute, including signing our Contributor License Agreement (CLA) and code review processes.</p>\n<h2>Star History</h2>\n<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://www.star-history.com/?repos=google-pytorch%2Ftorch_tpu&amp;type=date&amp;legend=top-left\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://camo.githubusercontent.com/5dd8f8b127ebb9d4b8e81870ade5b9dea014587bc270e5ff7c4984ca8bf88d43/68747470733a2f2f6170692e737461722d686973746f72792e636f6d2f63686172743f7265706f733d676f6f676c652d7079746f7263682f746f7263685f74707526747970653d64617465266c6567656e643d746f702d6c656674\" alt=\"Star History Chart\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<h2>License</h2>\n<p>This project is licensed under the Apache 2.0 License. See the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/google-pytorch/torch_tpu/blob/main/LICENSE\">LICENSE</a> file for details.</p>\n",
  "content_markdown": "[![TorchTPU Logo](https://github.com/google-pytorch/torch_tpu/raw/main/docs/images/torchtpu_logo.avif)](https://github.com/google-pytorch/torch_tpu/blob/main/docs/images/torchtpu_logo.avif)\n\n## Description\n\nThis 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\n## Features\n\n- PyTorch ATen kernel implementation for TPU.\n- Compilation cache for optimized execution of XLA computations.\n- Support for distributed training (Data Parallelism, Tensor Parallelism).\n- Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).\n- Integration with `torch.compile()` for graph mode execution.\n- Utilities for debugging and benchmarking.\n\n## User Guide and Interactive Notebooks\n\nA user guide is available at [https://torchtpu.readthedocs.io/](https://torchtpu.readthedocs.io/). This is a static version of the interactive notebooks. They are available in the docs folder.\n\nThe 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\n## Installation\n\nThe same instructions are available in the user guide linked above.\n\nTorchTPU's Python version support policy matches [PyTorch's Release Compatibility Matrix](https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix). Currently, **Python 3.11+** is supported. See the [Python Version Support Policy](https://github.com/google-pytorch/torch_tpu/blob/main/docs/developers/python_support_policy.md) for more details.\n\nFirst create a Python virtual environment:\n\n```shell\npython -m venv .venv\nsource .venv/bin/activate\n```\n\nor, with `uv`\n\n```shell\nuv venv .venv\nsource .venv/bin/activate\n```\n\n### Install via pip\n\nInstall the latest public release ([0.2.0 on PyPI](https://pypi.org/project/torch-tpu/0.2.0/)) (for `uv` prepend `uv` to all `pip` commands):\n\n```shell\npip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu\n```\n\n### Build from source\n\nWheels can be built via:\n\n```shell\nbazel build --config=wheel_common //ci/wheel:torch_tpu_wheel\n```\n\nAll 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\n```shell\nbazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \\\n    --//:allow_unoptimized_wheel=True\n```\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\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\n```shell\nbazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe\n```\n\nThe 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\nInstall wheel via:\n\n```shell\ncd <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\nNote 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\n## Getting Started\n\nSome pointers to get you started:\n\n- [PyTorch: Official PyTorch documentation](https://github.com/pytorch/pytorch)\n- [Tutorials and examples: easy to understand PyTorch code using TorchTPU](https://torchtpu.readthedocs.io/)\n- [Public APIs: the TorchTPU APIs that you can use](https://github.com/google-pytorch/torch_tpu/blob/main/docs/public_apis.md)\n\n## Dependencies\n\nThe project uses Bazel for dependency management. Key Python dependencies are:\n\nTBD\n\n## Bazel setup\n\nThis project uses Bazel for building and dependency management.\n\nTo install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.\n\nThere are several ways to install Bazelisk. Here are some common ways:\n\n- If you have Go installed:\n\n  ```\n  go install github.com/bazelbuild/bazelisk@latest\n  ```\n\n- Download the binary directly:\n\n  ```\n  wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64\n  ```\n\nSee the [Bazelisk README](https://github.com/bazelbuild/bazelisk/blob/master/README.md) for more install instructions\n\nHowever you install it, add it's location to `PATH`.\n\nTo verify the install, run `bazel info` in the repo's root directory:\n\n```\nbazel info\n```\n\nIt should print information about the repo without errors.\n\n## VSCode `clangd` Setup\n\nFor C++ code navigation, it is recommended to use VSCode with [`clangd` extension](https://marketplace.visualstudio.com/items?itemName=llvm-vs-code-extensions.vscode-clangd). Relying on a [compilation database](https://clang.llvm.org/docs/JSONCompilationDatabase.html) generated from Bazel action graph, [clangd](https://clangd.llvm.org/) is able to enrich the editor with various smart features, including code completions, go-to definitions, *etc.*\n\n1. Install VSCode extensions via UI or the following command:\n\n   ```shell\n   code --install-extension llvm-vs-code-extensions.vscode-clangd\n   # Uninstall Microsoft C++ extensions to avoid interference.\n   code --uninstall-extension ms-vscode.cpptools\n   ```\n\n2. Generate the initial `compile_commands.json`:\n\n   ```shell\n   tools/setup_clangd.py\n   ```\n\nIt 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](https://github.com/google-pytorch/torch_tpu/blob/main/tools/setup_clangd.py) file-level comments.\n\nAlternatively, you can invoke the script via `nox` (arguments after `--` are forwarded):\n\n```shell\nnox -s refresh_compile_commands -- --no-build\nnox -s refresh_compile_commands -- //src/torch_tpu/ops/...\n```\n\n## Contributing\n\nWe welcome contributions to the PyTorch on TPU project! Please see our [CONTRIBUTING.md](https://github.com/google-pytorch/torch_tpu/blob/main/CONTRIBUTING.md) for guidelines on how to contribute, including signing our Contributor License Agreement (CLA) and code review processes.\n\n## Star History\n\n[![Star History Chart](https://camo.githubusercontent.com/5dd8f8b127ebb9d4b8e81870ade5b9dea014587bc270e5ff7c4984ca8bf88d43/68747470733a2f2f6170692e737461722d686973746f72792e636f6d2f63686172743f7265706f733d676f6f676c652d7079746f7263682f746f7263685f74707526747970653d64617465266c6567656e643d746f702d6c656674)](https://www.star-history.com/?repos=google-pytorch%2Ftorch_tpu&type=date&legend=top-left)\n\n## License\n\nThis project is licensed under the Apache 2.0 License. See the [LICENSE](https://github.com/google-pytorch/torch_tpu/blob/main/LICENSE) file for details.\n",
  "dateCreated": "2026-10-20T21:31:53.761078291+00:00",
  "date_created": "2026-10-20T16:40:24+00:00",
  "date_modified": "2026-10-20T16:40:24+00:00",
  "date_published": "2026-10-20T16:40:24+00:00",
  "description": "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…",
  "external_url": "https://github.com/google-pytorch/torch_tpu",
  "git": {
    "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e",
    "path": "items/hnrss-org/2026/10/2026-10-20-google-open-sources-torchtpu-a-native-pytorch-back-end-for.md",
    "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/hnrss-org/2026/10/2026-10-20-google-open-sources-torchtpu-a-native-pytorch-back-end-for.md"
  },
  "id": "urn:aggr:item:ec897b5058ffe73a88e095b3f319f1a09c436003",
  "isBasedOn": {
    "@id": "https://github.com/google-pytorch/torch_tpu"
  },
  "mainEntity": {
    "@id": "https://github.com/google-pytorch/torch_tpu",
    "@type": "CreativeWork",
    "archivedAt": {
      "@id": "https://aggr.aymericbeaumet.com/items/hnrss-org/2026-10-20-google-open-sources-torchtpu-a-native-pytorch-back-end-for/#webpage"
    },
    "author": [
      {
        "@type": "Person",
        "name": "ilreb"
      }
    ],
    "dateModified": "2026-10-20T16:40:24+00:00",
    "datePublished": "2026-10-20T16:40:24+00:00",
    "headline": "Google open-sources TorchTPU, a native PyTorch back end for TPUs",
    "image": {
      "@type": "ImageObject",
      "caption": "Contribute to google-pytorch/torch_tpu development by creating an account on GitHub.",
      "height": 90,
      "url": "https://opengraph.githubassets.com/cd86c84e93f7b104a4d7723deb52a1b58d9418986b422db17cdc92ee64b72b49/google-pytorch/torch_tpu",
      "width": 160
    },
    "timeRequired": "PT5M",
    "url": "https://github.com/google-pytorch/torch_tpu",
    "wordCount": 909
  },
  "name": "Archived snapshot: Google open-sources TorchTPU, a native PyTorch back end for TPUs",
  "preview": {
    "alt": "Contribute to google-pytorch/torch_tpu development by creating an account on GitHub.",
    "color": null,
    "height": 90,
    "placeholder": {
      "data_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=",
      "hash": ""
    },
    "url": "https://opengraph.githubassets.com/cd86c84e93f7b104a4d7723deb52a1b58d9418986b422db17cdc92ee64b72b49/google-pytorch/torch_tpu",
    "width": 160
  },
  "reading_minutes": 5,
  "replicated_at": null,
  "source": "hnrss-org",
  "source_name": "Hacker News: Front Page",
  "summary": null,
  "tags": [],
  "temporalCoverage": "2026-10-20T21:31:53.761078291+00:00",
  "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/",
  "word_count": 909
}