Registry / ai-ml / torch-xla

torch-xla

JSON →
library2.9.0pypypiunverified

PyTorch/XLA is a Python package that bridges PyTorch with XLA devices (TPU, GPU, CPU) to enable high-performance machine learning. The current stable version is 2.9.0, with releases aligned to PyTorch minor versions. It supports Python 3.10-3.13 and provides both PJRT and XRT runtimes (PJRT recommended).

ai-mlawsgcp
pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
github ↗package ↗
INSTALL
IMPORT
SIG · TORCH-XLA
T
torch-xla
ai-mlpythonv2.9.0
Install
87.2s avg
Import
—
Disk
5555MB
Pass rate
4/ 10
Env Coverage4 / 10
glibc
3.9–3.13
musl
3.9–3.13
Install & Compatibility
Where this runs
tested against v2.9.0 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
glibc
py 3.10
✕ build_error
✓ 101s
py 3.11
✕ build_error
✓ 85s
py 3.12
✕ build_error
✓ 92.6s
py 3.13
✕ build_error
✓ 70.1s
py 3.9
✕ build_error
✕ timeout
5555MB installed
● package 5555MB
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

torch_xla
✓ import torch_xla
✗ from torch_xla import xla_model
torch_xla is a package; importing directly works. 'from torch_xla import xla_model' is outdated and may miss other submodules.
torch_xla.core.xla_model
✓ import torch_xla.core.xla_model as xm
✗ import torch_xla.xla_model as xm
The correct module path is 'torch_xla.core.xla_model', not a top-level xla_model.

Basic example: get XLA device, create tensor, run ops, mark step.

import torch import torch_xla import torch_xla.core.xla_model as xm # Get XLA device device = xm.xla_device() # Create tensor on XLA device t = torch.randn(3, 3, device=device) print(f"Tensor device: {t.device}") # Perform operations result = t + t print(f"Result: {result}") # Mark step and synchronize (required for XLA) xm.mark_step() xm.wait_device_ops()
Debug
Known issues
breakingXRT runtime is deprecated and removed in PyTorch/XLA 2.1+. Use PJRT runtime for all new code.
fix
Ensure you are using PJRT (default). If you explicitly used XRT, switch to PJRT by not setting XRT runtime env vars.
affects: >=2.1.0
gotchaMissing call to xm.mark_step() or xm.wait_device_ops() causes lazy execution to not materialize, leading to hangs or incorrect results.
fix
Call xm.mark_step() after each training step and xm.wait_device_ops() before measuring time or synchronizing.
affects: all
deprecatedLibrary installation extra [tpu] installs torch-xla but may conflict with newer versions of libtpu. Use separate pip install with index URL.
fix
Install libtpu directly: pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html for PyTorch/XLA 2.9+ it installs libtpu automatically.
affects: <2.9.0
Upgrade
Version history
2.9.0latest on PyPI · released Nov 17, 2025
Audit
Dependencies
torchrequiredPyTorch/XLA is an XLA bridge for PyTorch
jaxoptionalExperimental JAX interoperability
Agent activity
43 hits · last 30 days
node
38
OpenAI (training)
1
Resources
githubhttps://github.com/pytorch/xla ↗
packagehttps://pypi.org/project/torch-xla/ ↗
torch-xla — pip install torch-xla · libregistry