SOURCE-LINKED INTELLIGENCE

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

arXiv · AI, language, vision and robotics · article · Sep 16, 2026 · UTC

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 samp

Read original source ↗ Open in workspace

recordType

paper

region

Global