PyTorch to JAX: a field guide from porting a 22B video model

Jul 8, 2026
TPU

A field guide from converting LTX-2.3, a 22B audio-video generation model, across torch_xla, torchax, and native JAX on TPU7x.

Introduction

We recently converted LTX-2.3, Lightricks' open audio-video generation model, to run on Google's newest TPU generation (TPU7x, Ironwood). We tried three ways to do that: torch_xla, torchax, and a rewrite in JAX through MaxDiffusion.^[1] (torch_xla is PyTorch's native XLA backend; torchax is PyTorch running on JAX; MaxDiffusion is Google's native JAX diffusion implementation.)

We want to share our experience of bringing up a large real-world model on it, especially a model of this caliber: with 22B parameters across two modalities, plus VAEs, a text encoder, and a LoRA ecosystem, it exercises far more surface than a tutorial example.^[2] (Ironwood reached general availability two months before this project, and access to it is still limited.)

By the end, videos were coming out of both the original PyTorch implementation (ported via torchax) and the JAX implementation (using MaxDiffusion). The torchax port can run the full LTX-2.3 feature set on TPU, while MaxDiffusion runs its text-to-video natively and, after a contribution we merged upstream, also loads every official LTX-2.3 LoRA.

Image-to-video with synchronized generated audio, on a single TPU7x.

This post walks through the three routes in the order we used them, which is also the order of distance from PyTorch: torch_xla is closer to PyTorch and changes less code. torch_xla is PyTorch's own backend and changes the least code, torchax still lets you write PyTorch but runs it on JAX, and JAX is the full move. Everything below is about inference, on a single TPU host.

The test subject

LTX-2.3 generates video with synchronized audio from a text prompt, an image, or another video. The transformer holds 22B parameters in two parallel streams (a 14B video stream and a 5B audio stream, 48 blocks), and a Gemma-3-12B encoder processes the prompt. Altogether, roughly 70 GB of weights load before the first frame is decoded.

The public LTX codebase is wisely engineered to run within tight GPU memory budgets, which makes it a demanding conversion subject: the code works hard to fit the hardware, leaning on assumptions about CUDA, memory pressure, and eager execution at every level.

A single TPU7x chip, by contrast, has 192 GB of HBM, and everything it runs is compiled through XLA as whole programs. Where the GPU code juggles components in and out of memory, a TPU wants to load everything once, compile once, and reuse. That mismatch is the biggest source of porting work, but not the only one. The others we hit: semantic differences between the frameworks (mutability, random numbers, precision defaults), each codebase being tuned hard for its own hardware, and, on a chip this newly released, the fragility of the software stack itself, where every layer is running at its bleeding edge.

Route 1: torch_xla

The technique

torch_xla is PyTorch's native XLA backend and the longest-standing way to run PyTorch on TPUs: xla registers as a torch device, tensors placed on it trace their operations lazily into XLA graphs, which compile into an optimized program for the chip.

Running a PyTorch model with it requires minimal changes: request the XLA device (xm.xla_device()), move the model and its inputs to it, and add a sync point (xm.mark_step()) where results must materialize. Execution is lazy: operations accumulate into a graph, compilation happens at the sync point, and the model code itself stays untouched.

We tested this method on LTX-2's building blocks (the feed-forward, the SDPA attention, and a full audio-video transformer block) using a TPU v6e (Trillium), and it worked as documented: a stable install (torch 2.8 + torch_xla 2.8), with no code changes.

However, on the TPU7x, the same recipe failed for three different reasons, each failure hiding the next:

Stable torch_xla installs but fails at execute time with Unexpected PJRT_ExecuteOptions size: expected 112, got 80: its PJRT layer (the runtime interface between torch_xla and libtpu) predates every Ironwood-capable libtpu, so a nightly build is required.
The torch_xla[tpu] extra pins a libtpu that predates the chip, failing with TPU initialization failed: No ... device found. A newer libtpu had to be pinned over it.
On a bare VM, the instance metadata exposes TPU_-prefixed keys where torch_xla expects unprefixed ones, and startup dies with KeyError: 'ACCELERATOR_TYPE'. Bypassed with TPU_SKIP_MDS_QUERY=1 plus hand-set topology variables.

With all three patched, real transformer blocks computed on the chip, so the route can be made to work. Every fix, though, was found by digging from a stack trace down into version compatibility. Checking with Google, we discovered that this is not the recommended path for TPU7x: they pointed us to torchax (TorchTPU, a new native framework, is also on the way). So while this route remains viable on v6-generation TPUs, and probably on TPU7x as well, we wanted to convert in a way that will work with the newest hardware and without the fragile version stack.

Verdict. Keep your code: yes. Setup cost: low, but fragile on new chips. Result: transformer blocks verified on TPU v6e and TPU7x, not taken further. Fits when: you already run torch_xla on established TPUs. On new chips, use torchax or wait for TorchTPU.^[3] (Full recipe available in this gist.)

Route 2: torchax

The technique

torchax (formerly torch_xla2) is a PyTorch backend that provides graph-level interoperability between PyTorch and JAX. It means the nn.Module executes as-is, with tensors backed by JAX arrays, so XLA compilation, JAX sharding, and support for new TPU generations all come from the JAX side. And because the TPU support comes from JAX, setup on Ironwood required none of Route 1's workarounds: install, create the device, run.

In theory, running a PyTorch model with it requires minimal changes:

run torchax.enable_globally() and torch ops on device='jax' will lower through JAX.
convert the model to a JIT model using jitted = torchax.interop.JittableModule(model).

In practice, those changes covered 90% of the model. Another 5% failed loudly, and those failures were easy to fix. The last 5% failed silently: the code runs and the numbers are wrong. The rest of this route's story follows that split: the structural changes first, then the failures.

Inverting the memory model

Even on a large GPU, LTX-2 avoids holding all ~70 GB of weights resident at once: the model is built on the meta device, each component moves to the GPU right before use and is evicted afterwards.

On a TPU, however, the problem this machinery solves no longer exists: a TPU7x chip has 192 GB of HBM and the whole model fits. Worse, the machinery now works against the platform, because per-stage eviction fights the compile-once execution model: a compiled program pays off through reuse, called many times with its weights already resident, while eviction pays the host-to-device transfer again on every generation. So instead of porting the machinery, we replaced it with an EagerModelBuilder: build the final state dict on CPU (base weights with any LoRA deltas merged in, then cast to bfloat16), stream it tensor by tensor onto the device (avoiding two live copies of any weight), and keep the model resident across prompts.

To simplify the conversion and keep the port incremental, we also split the pipeline between CPU and TPU. The Gemma text encoder and audio decoder stayed on CPU, while the 22B transformer and the video VAE run on the TPU, with tensors moved across the boundary.^[4] (The Gemma text encoder does not trace under torchax because of its masking logic, and porting the audio decoder was not a priority.)

One trace of memory pressure survives even with 192 GB: at full resolution, the two-stage pipeline builds the stage-2 transformer as a fresh full copy (base + fused distilled LoRA) only after freeing stage 1. A reusable implementation should instead patch the resident one in place.

Silent failures

One failure was completely silent: every TPU run produced a broken clip, in eager and jit alike, while the same code was correct on an H200.

Same code and weights that render correctly on an H200. This is image-to-video, so the clip opens on the input frame and then degrades over its length into a washed-out, blocky field. The only difference from a correct run is whether the in-place writes landed.

It does not look like crisp TV static, because the garbage lives in the latent, and the VAE that decodes it can only emit its own smooth structure: the frames come out as a washed-out, blocky field. The cause was in the rotary position embedding:

# Build output, then patch its halves in place.
output = split_input * cos_freqs.unsqueeze(-2)
first_half_output = output[..., :1, :]
second_half_output = output[..., 1:, :]
first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)

This code hides a sneaky bug: in PyTorch, a slice is a view, so addcmul_ on the slice writes through into output. Under torchax, however, a slice of a JAX array is a copy. The addcmul_ calls therefore ran without error, modified two temporaries, and the results were discarded. output stayed at split_input * cos with no rotation applied, so Q and K carried no positional information, and the denoise collapsed into garbage.

The fix builds the halves out of place, with identical math:

cos = cos_freqs.unsqueeze(-2)
sin = sin_freqs.unsqueeze(-2)
first_half_output = first_half_input * cos - sin * second_half_input
second_half_output = second_half_input * cos + sin * first_half_input
output = torch.cat([first_half_output, second_half_output], dim=-2)

Once the mechanism was clear, we stopped debugging outputs and grepped the codebase for other in-place-on-view sites instead. There was exactly one more, in a VAE path we do not use.

Loud failures

These crashed with clear errors and cost far less time. Listed for completeness:

If the jitted function refers to the model directly instead of taking it as an argument, jit bakes its tens of GB of weights into the graph as constants and OOMs the compile. Pass the weights in as arguments instead (f(params, x), not f(x)) so they trace as inputs; torchax's JittableModule does this.
Custom dataclasses aren't JAX pytrees, so jax.jit errors on them.
A leftover @property alias on two conv wrappers broke torchax's parameter collection, which discovers params via dir().
torch.cuda.synchronize() in cleanup paths crashes a CPU-only torch build.
Enabling torchax globally intercepts CPU-side work too, including safetensors loading. CPU phases run under disable_temporarily().
Tensors created under @torch.inference_mode() error during tracing. so we run inference under torch.no_grad() instead.

Expected numerical differences

Two more behaviors diverge from CUDA, by design. They matter because until they are controlled, expected drift and real breakage look the same.

Different precision defaults. First, XLA reduces RMSNorm, LayerNorm, and softmax in bf16, where PyTorch on CUDA runs those reductions in fp32. Second, XLA computes fp32 matmuls at bf16 precision by default, rounding the fp32 operands to bf16 on the MXU, where CUDA keeps them in fp32 (jax_default_matmul_precision controls this). Neither raises an error, the output just drifts away from the CUDA reference. For validation against CUDA, one switch forces fp32 on both: explicit fp32 upcasts around every reduction, plus jax_default_matmul_precision="highest" for the matmuls. Production keeps the fast bf16 defaults.

A different random number generator. Under torchax, torch.randn draws from JAX's RNG. Seeds still work: different seeds give different noise, and a given seed is reproducible on the TPU side, but the stream comes from a different algorithm, so the same seed produces completely different numbers than PyTorch on CUDA (and the seeded torch.Generator our pipeline passed was ignored outright).

Debugging method

Silent failures leave no stack trace, so finding them meant comparing the TPU run against the CUDA reference numerically, at matched points along the pipeline. The comparison is only meaningful once the expected differences above are pinned: fixed noise for the RNG, and fp32 mode for precision. To do that, we built a harness that:

Recorded ground truth on CUDA: Patched the noiser to save each noise tensor it draws, and save the final latent before VAE decode.
Replayed on TPU: Wrote a FixedNoiser injects the recorded noise instead of drawing fresh.
And compared the latents (max-abs, mean-abs), expecting a small stable diff.

Beyond the harness, each fix also landed with a smoke test on the CPU JAX backend (JAX_PLATFORMS=cpu), which needs no TPU and no weights. The result is that the conversion layer, the part most likely to regress, is covered by tests that run anywhere.

Results

The port runs on TPU with the full LTX-2.3 feature set: the two-stage pipeline (base denoise, latent upsample, distilled refine), image conditioning, video-to-video, and IC-LoRA reference conditioning. For example, the colorization video below was produced with it: greyscale input, IC-LoRA fused, color output, end to end on a TPU7x.

Greyscale input next to the colorized output. The colorized output carries synchronized generated audio.

Performance-wise, the first run at a given shape (the prompt encoder pads to a fixed token length) is dominated by XLA compilation, roughly 2.5x the warm time, and JAX's persistent compilation cache carries the compiled programs across process restarts (at this scale jit compiling is practically a must - eager execution through torchax costs tens of seconds per transformer forward). Warm runs land in the same ballpark as a high-end GPU, before any TPU-specific optimization. Full benchmarks will be released in the future.

The code is public: moonmath-ai/LTX-2-with-torchax.

Verdict. Keep your code: yes. Setup cost: near zero. Porting cost: low-medium. Result: full feature parity on TPU. Fits when: the PyTorch implementation is the source of truth and keeps moving.

Route 3: a native JAX implementation (MaxDiffusion)

The technique

An automatic conversion from PyTorch to JAX can hit a ceiling, so the third direction is a full JAX rewrite. Fortunately, we did not have to write one: Google ships MaxDiffusion and MaxText, reference implementations of common diffusion and language models, written with JAX-native idioms. MaxDiffusion already includes an LTX-2 pipeline, TPU-first: sharding the 22B across cores is a mesh flag, not a code change, and the attention kernels are TPU-native.

The cost of reimplementation is the usual: it’s a separate codebase, it takes time to nail right, and it lags its reference. In our case, the LTX-2 pipeline in MaxDiffusion is text-to-video only: no image conditioning, no reference-video conditioning, and no second-stage refinement.

Getting it running was its own small bring-up. Most of it was just setting the right values: the single-host distribution flag and a sharding mesh sized to the chip's core count. Two things did need care: The config's default attention kernel crashes the audio stream on 2.3, so you have to select a2v_attention_kernel=dot_product, and MaxDiffusion loads only the Diffusers weight layout from a community re-upload in the right layout. We did hit several segfaults while running in various modes, and opened issues upstream for them.

Text-to-video with synchronized generated audio from MaxDiffusion on a TPU7x: 1024x1536, 121 frames, single stage. Different prompt and configuration than the clips above, so not a side-by-side comparison.

What we contributed

While bringing LTX-2.3 up on MaxDiffusion, we found one thing genuinely broken and fixed it upstream: LTX-2.3 (22B) LoRA checkpoints did not fully load. The root of it was naming: LoRA checkpoints name their weights after the official LTX modules, while MaxDiffusion loads its base weights from a Diffusers-layout repo with different module names, so a key map translates between the two namespaces. That map predated some of the 22B's layer families, and the loader skipped the unmatched keys silently, so a partially-mapped LoRA fused into a partially-adapted model with no warning.

We extended the translation map, changed the diagnostics to report unmatched keys instead of dropping them, and verified every official LTX-2.3 LoRA against the real model. The change is merged into MaxDiffusion main.

The remaining gaps

While the LoRAs load correctly, some of the full LTX feature set is still missing from MaxDiffusion's pipeline, though most of the machinery already exists:

Image- and video-conditioned generation. The pipeline already accepts an initial latent with a noise scale (SDEdit-style), so wiring in a VAE-encoded image or video is tens of lines.
Pipelines that exist only in the LTX repo, for example two-stage refinement. Re-noising, custom sigma schedules, and no-guidance denoising all exist, and the distilled LoRA the second stage requires now loads. What remains is orchestration.
IC-LoRA reference conditioning (what the colorization demo uses) is the real gap: reference frames concatenated into the token sequence with a conditioning mask, roughly a few hundred lines. Until it exists, IC-LoRAs load but have no input to condition on.

Verdict. Keep your code: no. Setup cost: low if the model exists, medium to add features, high if it needs a full rewrite. Result: native text-to-video on TPU7x, LoRA support merged upstream. Fits when: TPU-first production and the feature set covers your use case.

How to choose

Each route is the path of least resistance for a different starting point.

If you already run torch_xla on established TPUs, staying there is the cheapest option: it works, and the pain is concentrated at new-chip launches. If you have a living PyTorch codebase, torchax is the cheapest conversion, with the caveat that the last 5% could fail silently, so budget for the validation harness and not just the port. If you are TPU-first and want the highest performance ceiling, a native JAX implementation gives the compiler, the sharding, and the kernels first-class treatment, at the cost of a separate codebase to maintain or adopt.

	torch_xla	torchax port	MaxDiffusion
Status on TPU7x today:	proof of concept	working, full feature set	working, text-to-video (+ our LoRA fix)
Keep your PyTorch code:	yes	yes	no
Brand-new hardware:	fragile version stack	clean (rides JAX)	clean (is JAX)
What broke for us:	loud: ABI and environment errors	silent: dropped writes, numerics	silent: checkpoint key mapping
Performance notes:	lazy tracing, can be host-bound	same HLO primitives as JAX*	adds hand-tuned TPU kernels and direct sharding control
Pick it when:	already invested, established chips	living PyTorch codebase, need every feature	TPU-first production, features covered

The failure row records what we hit, not everything each route can produce. In particular, we did not take torch_xla far enough to meet its silent class (unexpected recompiles and device syncs, which degrade performance without changing output).

For us the two working routes turned out to be complementary rather than competing. The port gave us the full feature set and a reference implementation running on the same hardware. The native implementation gave us a performance baseline and the codebase where the ecosystem converges. The implementation knowledge moved from the port to the upstream contribution, and the upstream baseline told us where the port's performance should land.

What we learned

Porting lessons

The expensive failures are silent, so build the comparison harness early. Dropped in-place writes, mismatched precision defaults, a different RNG, loaders that skip unmatched keys: each produced a running model whose output was wrong or subtly different, with no error pointing anywhere. What caught them was the comparison harness: record the CUDA run's noise and latents, replay them on TPU, and diff the numbers.

Cutting-edge hardware means a cutting-edge stack. We started about two months after TPU7x went GA, and at that point the chip was supported only by the newest layer of each library, and not by all of them. Finding a stack that knows the chip exists is part of the work, and it may mean switching stacks, as it did for us. It also means being ready for faults, and ready to fix them: a version mismatch can surface three layers down, and a young stack can simply segfault.

What TPU and JAX buy you

Whole-program compilation. XLA compiles the full denoise step as one program and optimizes across operation boundaries. Where eager PyTorch dispatches each operation as its own kernel, round-tripping intermediates through memory each time, XLA fuses long chains of ops into a handful of kernels, which cuts memory-bandwidth traffic and keeps intermediates in fast on-chip memory. TPUs are built for exactly this: a large systolic-array matmul unit that pays off when the compiler can hand it big fused regions instead of op-at-a-time work. The catch is that XLA specializes to the shapes it sees, so a new shape triggers a fresh compile. The compile is one time per shape (compiled programs can persist across process restarts through JAX's compilation cache), so we aim to use pipelines with static shapes.^[5] (PyTorch's own torch.compile brings similar fusion to eager code, but it captures partial graphs and falls back to eager on anything it cannot trace. XLA compiles the region as one whole program with no eager fallback.)

Sharding as configuration. JAX describes the devices as a named mesh and lets you annotate which array axes to split across it. XLA's SPMD partitioner (GSPMD) then rewrites the single-device program into a distributed one, inserting the all-gathers and reduce-scatters itself. It is the same partitioning machinery large-model training uses in MaxText, so data-, tensor-, and context-parallelism become config axes rather than code rewrites. Sharding the 22B across cores is a mesh axis in a config file, and the compiler works out the communication.^[6] (On GPUs this is typically a library wrapping the model, Megatron-LM or DeepSpeed or PyTorch FSDP, rather than compiler partitioning from annotations. PyTorch's newer DTensor is converging on the same model, and GSPMD also runs on JAX-on-GPU.)

A hardware-free development loop. JAX exposes one NumPy-like API that lowers to XLA for whichever backend is selected, so the same source runs on CPU, GPU, or TPU by changing a single environment variable. That makes the CPU backend a faithful stand-in for logic and shape checking, minus performance. The same code runs on JAX's CPU backend, so every fix in the port landed with a smoke test that needs no TPU and no weights. The correctness of the conversion layer is testable on a laptop, and in CI.

What we shipped

The work produced two deliverables we are happy to share. Our torchax port, moonmath-ai/LTX-2-with-torchax, runs the full LTX-2.3 feature set on TPU7x, including the two-stage pipeline and IC-LoRA video-to-video. And MaxDiffusion, as of our merged contribution, loads every official LTX-2.3 LoRA cleanly on top of its native text-to-video pipeline. The remaining MaxDiffusion gaps (conditioning, two-stage) are closable with machinery that mostly exists. TorchTPU, Google's native PyTorch framework for TPUs, was announced in April 2026 and is already running internally at Google, with a public release planned for later in 2026. It is an exciting addition that may change this map again.

Thanks to Google's TPU team for the guidance and to Lightricks for the open model.

TL;DR

Three routes: torch_xla, torchax, and native JAX through MaxDiffusion.
The full torchax port reaches LTX-2.3 feature parity on TPU7x.
MaxDiffusion now loads every official LTX-2.3 LoRA after the upstream contribution.

torch_xla is PyTorch's native XLA backend; torchax is PyTorch running on JAX; MaxDiffusion is Google's native JAX diffusion implementation.
Ironwood reached general availability two months before this project, and access to it is still limited.
Full recipe available in this gist.
The Gemma text encoder does not trace under torchax because of its masking logic, and porting the audio decoder was not a priority.
PyTorch's own torch.compile brings similar fusion to eager code, but it captures partial graphs and falls back to eager on anything it cannot trace. XLA compiles the region as one whole program with no eager fallback.
On GPUs this is typically a library wrapping the model, Megatron-LM or DeepSpeed or PyTorch FSDP, rather than compiler partitioning from annotations. PyTorch's newer DTensor is converging on the same model, and GSPMD also runs on JAX-on-GPU.