JAX
More info
- Creator
- Released
- Owner
- Type
- ML framework
- License
- Apache 2.0
- Website
- jax.readthedocs.io
- Source
- github.com/jax-ml/jax
- Wikipedia
- en.wikipedia.org
Google's ML framework — NumPy API plus automatic differentiation, JIT compilation via XLA, and first-class TPU support.
Functional paradigm: no mutable state, models are functions from (params, x) → y. Four magical transforms: jit (compile), grad (autodiff), vmap (auto-batching), pmap/shard_map (multi-device). Allows extremely compact research code.
Drives Google DeepMind (Gemini, AlphaFold), Anthropic's Claude training, and a large slice of Google's internal ML. PyTorch dominates the industry, but JAX is almost mandatory if you're working on TPU.