JAX
Plus d'infos
- Créé par
- Année
- Propriétaire
- Type
- Framework ML
- Licence
- Apache 2.0
- Site officiel
- jax.readthedocs.io
- Code source
- github.com/jax-ml/jax
Le framework de ML de Google — une API façon NumPy plus la différentiation automatique, la compilation JIT via XLA, et un support de premier ordre pour les TPU.
Paradigme fonctionnel : aucun état mutable, les modèles sont des fonctions de (params, x) → y. Quatre transformations magiques : jit (compiler), grad (autodiff), vmap (batching automatique), pmap/shard_map (multi-device). Permet un code de recherche extrêmement compact.
JAX fait tourner Google DeepMind (Gemini, AlphaFold), l'entraînement de Claude chez Anthropic, et une bonne part du ML interne de Google. PyTorch domine l'industrie, mais JAX est quasiment incontournable pour qui travaille sur TPU.