JAX
Software
Training
Active
Composable function transforms — the framework Gemini is built with.
You're looking at a software record: a tool or platform in the AI stack — who maintains it and what hardware it runs on.
Photo: Google LLC · Public domain · Wikimedia Commons
JAX treats models as pure functions to be transformed — grad, jit, vmap, pmap — compiling through XLA to TPU pods with exceptional scaling ergonomics. It is Google DeepMind's house framework (Gemini training) and the favorite of much of the research frontier outside PyTorch.
How this connects to the rest of the stack
This record's links at a glance — click a group to jump to it.
Runs on
Maintained by
Competes with
Related
Sources
Record reviewed as of 2026-01-15 — see Methodology.