# vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes

arXiv: https://arxiv.org/abs/2609.18077 · HTML: https://mlsanity.com/paper/2609.18077/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 20 references that are in the map.
- First posted on arXiv 16 Sep 2026.

## Builds on (most cited first, 18 of 18)

- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer](https://mlsanity.com/paper/2408.06072/) (2024)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [HunyuanVideo: A Systematic Framework For Large Video Generative Models](https://mlsanity.com/paper/2412.03603/) (2024)
- [LTX-Video: Realtime Video Latent Diffusion](https://mlsanity.com/paper/2501.00103/) (2024)
- [UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models](https://mlsanity.com/paper/2302.04867/) (2023)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [LTX-2: Efficient Joint Audio-Visual Foundation Model](https://mlsanity.com/paper/2601.03233/) (2026)
- [World Simulation with Video Foundation Models for Physical AI](https://mlsanity.com/paper/2511.00062/) (2025)
- [Cosmos 3: Omnimodal World Models for Physical AI](https://mlsanity.com/paper/2606.02800/) (2026)
- [HunyuanVideo 1.5 Technical Report](https://mlsanity.com/paper/2511.18870/) (2025)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)

## Abstract (arXiv)

Open-source video generative models ship almost exclusively as PyTorch/CUDA reference implementations. This leaves Cloud TPU pods without a production-ready inference path, despite offering large, cost-effective accelerator memory pools ideal for long-sequence spatiotemporal attention. We present vidax, an open-source JAX/Flax inference engine and zero-copy PyTorch-to-JAX weight translator for modern video generation architectures. vidax covers a diverse set of spatiotemporal models --- including Diffusion Transformers, omnimodal Mixture-of-Transformers, 3D VAEs, text encoders, and native samplers --- with zero PyTorch dependency in the execution path. The framework unifies 1D tensor parallelism with DeepSpeed-Ulysses sequence parallelism on a single JAX sharding mesh, integrates TPU flash-attention kernels, and implements per-layer weight offloading to support reference resolutions that exceed single-device memory. We benchmark compile times, latency, and peak memory utilization on TPU v4-8 hardware, and document real-world numerical bugs surfaced during checkpoint translation. vidax is released open-source as a baseline for JAX and TPU video generation research.