Side project ports most video generation models from PyTorch to Jax for TPU

ceciletamura · x · 2026-09-24

Researcher CongyueD released a side-project TPU library that translates the majority of video generation models from PyTorch to Jax, enabling them to run on Google TPUs. The project ships with source code, a tech report, and a quick PyPI install for trying it out, built with support from Google's TPU Research Cloud (TRC). A ready-made path for researchers who want video models in the TPU ecosystem.

Original post →

More from Infra

Infra channel →