---
id: torchax
version: "0.0.13"
license: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# torchax — torchax is a library for running Jax and PyTorch together

License: permissive · Maintenance: active · Downloads: 164.6K/mo

## 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 above — 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

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

- License: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 164.6K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

pytorch on tpu, pytorch jax interop, run pytorch with jax, tpu backend for pytorch, pytorch jax bridge, jax array torch tensor, pytorch tpu support, tpu-backend, pytorch-jax-interop, accelerator-support

[View on SkillFed](https://skillfed.io/packages/torchax) · [View on PyPI](https://pypi.org/project/torchax/)