JAX
Mer info
- Skapad av
- År
- Ägare
- Typ
- ML-ramverk
- Licens
- Apache 2.0
- Officiell sajt
- jax.readthedocs.io
- Källkod
- github.com/jax-ml/jax
Googles ML-framework — NumPy-API plus automatisk differentiering, JIT-kompilering via XLA, och första-klass-stöd för TPU.
Funktionell paradigm: ingen muterbar state, modeller är funktioner från (params, x) → y. Fyra magiska transformer: jit (kompilera), grad (autodiff), vmap (auto-batching), pmap/shard_map (multi-device). Tillåter superkompakt research-kod.
Driver Google DeepMind (Gemini, AlphaFold), Anthropic Claude-träning, och en stor del av Googles interna ML. PyTorch dominerar i industrin men JAX är nästan obligatoriskt om du jobbar på TPU.