TorchAx: PyTorch on JAX/TPU Bridge

Warnings

• gotcha Enabling `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')`.
• gotcha Running `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.
• gotcha JAX'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.
• gotcha JAX 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.
• gotcha When 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.

Install

• pip install torchax Install torchax
• # First, install PyTorch CPU: pip install torch --index-url https://download.pytorch.org/whl/cpu # Linux pip install torch # Mac # Then, install JAX for your accelerator: pip install -U jax[tpu] # Google Cloud TPU pip install -U jax[cuda12] # GPU machines pip install -U jax # Linux CPU or Mac Prerequisites: PyTorch and JAX

Imports

• torchax

  import torchax

  The primary import for enabling torchax functionality.
• enable_globally

  torchax.enable_globally()

  Activates torchax to intercept PyTorch operations and execute them via JAX. Should be called after model loading.
• JittableModule

  from torchax.interop import JittableModule

  Helper class to wrap PyTorch models for JIT compilation with JAX.
• jax_jit

  from torchax.interop import jax_jit

  Decorator to compile PyTorch functions (taking/returning torch.Tensors) into JAX-compiled versions for performance.

Quickstart

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)