vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes
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 para...
| Authors | Congyue Deng |
|---|---|
| Primary category | cs.CV |
| Published | 16 September 2026 |
| Last updated at source | 16 September 2026 |
Silytics shows public metadata from arXiv and links to the original. Information may be incomplete or out of date. Last refreshed 17 October 2026. Licenses and terms of the original apply to the work itself. Not financial, legal or professional advice.