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

arXiv:2609.18077 · cs.CV, cs.DC, cs.LG · Submitted 2026-09-16 · Read on arXiv

cs.CV, cs.DC, cs.LG

Submitted: 2026-09-16

Updated: 2026-09-16

Code: https://github.com/FlyingGiraffe/vidax

License: http://creativecommons.org/licenses/by/4.0/

The gist: Open-source video generative models ship almost exclusively as PyTorch/CUDA reference implementations.

Terminology

Abstract

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.

Sources

Related papers