torchax

library0.0.13pypypi✓ verified 27d ago

torchax is a library that serves as a backend for PyTorch, enabling users to run PyTorch programs on JAX-supported hardware like Google Cloud TPUs. It provides seamless graph-level interoperability, allowing the mixing of JAX and PyTorch syntax within the same program, and leveraging JAX features such as `jax.grad`, Optax, and GSPMD for PyTorch model training. The current version is 0.0.11, with development active on GitHub.

ai-ml

pip install torchax

github ↗package ↗

INSTALL

IMPORT

SIG · TORCHAX

torchax

ai-mlpythonv0.0.13

Install

1.6s avg

Import

—

Disk

16MB

Pass rate

8/ 10

Env Coverage8 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.0.7 · 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

✓ —

✓ 1.6s

py 3.11

✓ —

✓ 1.7s

py 3.12

✓ —

✓ 1.6s

py 3.13

✓ —

✓ 1.6s

py 3.9

✕ build_error

16MB installed

● package 16MB

Code

Verified usage

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

torchax

✓ import torchax

✗ import torchax

This quickstart demonstrates how to run a standard PyTorch `nn.Module` using torchax. The key steps are to import `torchax`, call `torchax.enable_globally()` *after* model initialization, and then move the model and inputs to the 'jax' device. For improved performance, especially in production, `torchax.interop.JittableModule` (which leverages `jax.jit`) is recommended for compiling the model.

import torch import torch.nn as nn import torch.nn.functional as F import torchax class MyModel(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(28 * 28, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = x.view(-1, 28 * 28) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x # Instantiate the PyTorch model m = MyModel() # IMPORTANT: Enable torchax GLOBALLY *after* model instantiation/loading torchax.enable_globally() # Move the model to the 'jax' device m.to('jax') # Create input tensor on the 'jax' device inputs = torch.randn(3, 1, 28, 28, device='jax') # Run the model; operations will be executed by JAX outputs = m(inputs) print(outputs.shape) print(outputs.device) # Example with jax.jit for performance (using JittableModule) from torchax.interop import JittableModule m_jitted = JittableModule(m) # Wraps the model for JIT compilation jitted_outputs = m_jitted(inputs) print(jitted_outputs.shape) print(jitted_outputs.device)

Debug

Known issues

gotchaEnabling `torchax.enable_globally()` before loading a PyTorch model can lead to errors, as it might intercept unsupported initialization operations. Always enable globally *after* the model has been fully loaded or instantiated.

fix

Ensure `torchax.enable_globally()` is called only after your `torch.nn.Module` has been initialized and its weights potentially loaded. For example, `model = MyModel(); torchax.enable_globally(); model.to('jax')`.

affects: All versions

gotchaRunning `torchax` models in eager mode (without JAX JIT compilation) can be significantly slower than native PyTorch or JIT-compiled JAX execution. JAX's eager mode generally does not offer the same performance benefits as its compiled mode.

fix

For performance-critical workloads, always use JAX's Just-In-Time (JIT) compilation. Wrap your model with `torchax.interop.JittableModule` or decorate functions with `torchax.interop.jax_jit` to compile the computation graph for faster execution. The first call will include compilation time, but subsequent calls will be much faster.

affects: All versions

gotchaJAX's JIT compilation specializes for fixed input shapes. If input shapes change between calls (common in scenarios like autoregressive text generation), JAX will recompile the graph, potentially leading to performance degradation worse than eager mode.

fix

For dynamic input shapes, consider using techniques like `StaticCache` (for Hugging Face models) or ensuring that varying dimensions are handled as static arguments (`static_argnums`) if using `jax.jit` directly. Alternatively, refactor the computation to minimize shape changes within JIT-compiled regions.

affects: All versions

gotchaJAX transformations, including JIT, require functions to be 'pure' (i.e., all inputs passed as arguments, all outputs returned, no side effects or closure over mutable state). PyTorch `nn.Module.forward` implicitly closes over model weights. This can lead to unexpected behavior or performance issues with JAX.

fix

Utilize `torchax.interop.JittableModule` which handles this by passing weights as explicit arguments, or explicitly convert your PyTorch model to a functional form using `torch.func.functional_call` when interacting with JAX transforms.

affects: All versions

gotchaWhen interoperating with custom JAX types (e.g., specific output types from HuggingFace models like `CausalLMOutputWithPast`), these types might not be automatically recognized by JAX's pytree mechanism. This can cause `TypeError: ... is not a valid JAX type` errors.

fix

Register custom types as JAX pytrees using `jax.tree_util.register_pytree_node`. Refer to JAX documentation and `torchax` examples for correct registration patterns.

affects: All versions

Upgrade

Version history

0.0.13latest on PyPI · released Jun 17, 2026

Audit

Dependencies

torchrequiredtorchax is a PyTorch backend/frontend for JAX, requiring PyTorch to function. Users must choose their desired PyTorch build (CPU, CUDA, etc.)

jaxrequiredtorchax runs PyTorch models on JAX's backend, requiring JAX to be installed with the appropriate accelerator (TPU, CUDA, or CPU).

Agent activity

28 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/google/torchax ↗

packagehttps://pypi.org/project/torchax/ ↗