JAX
Altre info
- Creatore
- Anno
- Proprietario
- Tipo
- Framework di ML
- Licenza
- Apache 2.0
- Sito ufficiale
- jax.readthedocs.io
- Codice sorgente
- github.com/jax-ml/jax
- Wikipedia
- en.wikipedia.org
Il framework di ML di Google: API di NumPy, differenziazione automatica, compilazione JIT tramite XLA e supporto di prima classe per le TPU.
Paradigma funzionale: niente stato mutabile, i modelli sono funzioni da (parametri, x) → y. Quattro trasformazioni quasi magiche: jit (compilazione), grad (autodiff), vmap (batching automatico), pmap/shard_map (più dispositivi). Permette codice di ricerca estremamente compatto.
Alimenta Google DeepMind (Gemini, AlphaFold), l'addestramento di Claude di Anthropic e una buona fetta del ML interno di Google. Nell'industria domina PyTorch, ma con le TPU JAX è quasi obbligatorio.