Whisper JAX

Whisper JAX reimplements OpenAI's Whisper in JAX and Flax on top of the Hugging Face Transformers codebase, targeting raw transcription throughput on accelerator hardware. Its speed comes from JIT compilation, data parallelism across devices via JAX's pmap, optional float16 or bfloat16 computation, and a batching strategy that splits long audio into 30 second chunks transcribed in parallel rather than sequentially. Benchmarks in the repository show a single A100 GPU transcribing an hour of audio in about 75 seconds against roughly 1,000 seconds for the original PyTorch version, and a TPU v4-8 doing the same in under 14 seconds, the source of its headline claim of over 70 times faster inference. All Whisper checkpoints from tiny through large-v2 are supported for multilingual transcription, translation, and timestamps, and the code runs on CPU, GPU, or TPU, standalone or as a Gradio inference endpoint. Open source under the Apache 2.0 license, it suits researchers and services that batch-process large audio archives on TPU or GPU infrastructure.