JAX
Google/DeepMind's training framework, the PyTorch alternative
JAX is the functional autodiff framework behind Google DeepMind's Gemini and most TPU-first training; its pmap/pjit and Pathways runtime enable planet-scale TPU pods. It is the principal reason a viable non-PyTorch, non-CUDA training path exists at frontier scale.
Owner
Google DeepMind
Runs on
TPU via XLA
How it fits the stack
JAX with what it depends on (above) and what it feeds (below). The figure renders as a crawlable diagram and upgrades to an interactive 3D graph as it scrolls into view.
JAX in the AI stack. JAX with its immediate upstream dependencies (top) and downstream dependents (bottom) in the AI value chain. Hover a node in 3D, or read the full relationships below.
Graph data (text) — 5 entities, 5 relationships
- JAX —depends on→ Google TPU
- JAX —depends on→ MLIR / OpenXLA / StableHLO
- Gemini —uses→ JAX
- Google DeepMind —uses→ JAX
- JAX —designs→ Google DeepMind
Depends on ↑ · 2
Context — capital, rivals, policy · · 1