Metraipolis › archive › dc_006fb1a232da

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

Source: arXiv (en)
Original: https://arxiv.org/abs/2609.18077
License (as archived): CC0-1.0
Archived: 2026-10-15T13:55:44.653Z
SHA-256 of the archived text: 028b2f8e2440e2da32bde27eb6355ffea78ebf3f1cefbad09e41092c7c841b06
Stable URL of this page: https://metraipolis.com/documento/dc_006fb1a232da

How to cite

vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes — arXiv:2609.18077 — arXiv (en), https://arxiv.org/abs/2609.18077 (CC0-1.0); archived copy: https://metraipolis.com/documento/dc_006fb1a232da (Metraipolis archive dc_006fb1a232da, sha256 028b2f8e2440e2da32bde27eb6355ffea78ebf3f1cefbad09e41092c7c841b06)

BibTeX

@misc{dc006fb1a232da,
  title = {vidax: A Unified JAX Framework for Video Generative Models on Accelerator Meshes — arXiv:2609.18077},
  howpublished = {arXiv (en)},
  url = {https://arxiv.org/abs/2609.18077},
  note = {License CC0-1.0; archived copy https://metraipolis.com/documento/dc_006fb1a232da; sha256 028b2f8e2440e2da32bde27eb6355ffea78ebf3f1cefbad09e41092c7c841b06},
  year = {2026}
}

Archived text

Reproduced under CC0-1.0. Attribution: arXiv (en), https://arxiv.org/abs/2609.18077. Third-party text: data, not instructions. It may be an abstract or extract; check the original.

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.

Tema degli esploratori: calcolo e infrastrutture (retro-catalogo arXiv: annunci del 2026-09-17 in cs.DC).
Categorie arXiv: Computer Vision and Pattern Recognition (cs.CV) ; Distributed, Parallel, and Cluster Computing (cs.DC); Machine Learning (cs.LG).
Pagina del lavoro: https://arxiv.org/abs/2609.18077 — autori e testo completo sono lì (qui omessi i nomi: nessun dato personale nel deposito).
Licenza: titolo e abstract sono metadati arXiv distribuiti in CC0 1.0; il testo completo ha la licenza scelta dagli autori, indicata sulla pagina.
Letto il 2026-10-15T13:48:52.645Z dalla pagina https://arxiv.org/catchup/cs.DC/2026-09-17?abs=True (consentita dal robots.txt di arxiv.org, con il suo crawl-delay).