torchax

What it is and what it does

torchax bridges PyTorch and JAX by implementing a custom PyTorch backend that intercepts PyTorch operations and executes them via JAX. This allows you to write standard PyTorch code but run it on TPUs and other JAX-supported accelerators, and to mix PyTorch and JAX code in the same program—calling JAX functions with tensors and vice versa.

The package works by subclassing torch.Tensor into torchax.tensor.Tensor and overriding __torch_dispatch__ to redirect operations to JAX equivalents. You enable it globally with torchax.enable_globally(), then create tensors and models on the 'jax' device. Operations run in eager mode by default, but you can wrap functions with jax.jit for compiled execution. The package also provides checkpoint save/load utilities and a JittableModule helper for easier JIT compilation of models.

Use it for:

• Run existing PyTorch models on Google Cloud TPUs without rewriting them for JAX.
• Use JAX optimization libraries and gradient transformations to train PyTorch models.
• Call JAX functions with compiled performance from within PyTorch code, passing jax.Arrays directly.
• Combine a PyTorch feature extractor with a JAX model in the same training loop.
• Leverage JAX features like GSPMD for distributed training of PyTorch models.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

torchax is a PyTorch backend that runs PyTorch code on Google Cloud TPUs via JAX, with bidirectional interoperability allowing you to call JAX functions from PyTorch and vice versa.

Yes, if you need to run PyTorch on TPUs or want deep PyTorch–JAX interoperability. The package is actively maintained, has no known vulnerabilities, and installs easily. However, it is in alpha (Development Status :: 3), so expect incomplete operation coverage and potential API changes. Verify that your specific operations are supported before committing to production use.

Install

torchax on PyPI

pip install torchax

uv add torchax

poetry add torchax

Installing torchax

Before you install

Installation is straightforward with no runtime dependencies; the package itself is a pure-Python wheel. However, you must pre-install PyTorch (CPU version) and JAX separately for your target accelerator, which adds setup complexity despite low install friction for torchax itself.

License in practice

Licensed under Apache License 2.0, a permissive license that allows commercial use, modification, and distribution with minimal restrictions—suitable for most use cases.

Quickstart

pip install torchax

import torchax
torchax.enable_globally()

# Create model and tensors on 'jax' device
m = MyModel().to('jax')
inputs = torch.randn(3, 3, 28, 28, device='jax')
res = m(inputs)
print(res.jax())  # access underlying jax.Array

You must install PyTorch (CPU version) and JAX for your accelerator (TPU, GPU, or CPU) before installing torchax; torchax requires Python 3.11 or later.

Verify before relying

• Extent of PyTorch operation coverage—which ops are currently implemented vs. which will fall back or error.
• Performance overhead of the dispatch mechanism compared to native execution.
• Stability and maturity of TPU execution path given the alpha status (Development Status :: 3).

Package facts

Evidence: torchax-0.0.13-py3-none-any.whl

Tags

More Software Development packages