Introduzione a JAX e alla costruzione di modelli di machine learning 2026

Che cos'è JAX, come si costruisce e si addestra un modello, che cosa costa una TPU rispetto a una GPU, dove JAX non ha rivali e dove è la scelta sbagliata

JAX è il meno conosciuto e il più frainteso dei quattro strumenti di questa serie, e il fraintendimento comincia dalla categoria in cui lo si mette. JAX non è un framework di deep learning: la sua stessa documentazione lo definisce «una libreria Python per il calcolo su array orientato agli acceleratori e per la trasformazione di programmi». Non contiene strati di rete neurale, non ha un ciclo di addestramento, non ha un ottimizzatore. Contiene quattro o cinque trasformazioni componibili — derivare, vettorizzare, compilare, distribuire — che si applicano a normale codice numerico Python. Tutto il resto sta in librerie separate, per scelta dichiarata.

Questa guida percorre JAX per intero — caratteristiche principali, vantaggi competitivi, costi, funzionalità avanzate e integrazioni, alternative e confronto, strategie di crescita e tendenze — ed è la quarta di una serie dopo quelle su TensorFlow, PyTorch e scikit-learn. Arriva per ultima per una ragione precisa: è lo strumento che si sceglie quando si sa già perché gli altri tre non bastano. Se non avete quella ragione, questa guida vi servirà soprattutto a capire che non vi serve — il che vale comunque il tempo di lettura.

Due fatti la rendono però meno di nicchia di quanto sembri. Il primo: il 6 marzo 2026, mettendo TensorFlow in manutenzione, Google ha raccomandato «Keras 3, JAX e PyTorch per il nuovo lavoro sull'intelligenza artificiale generativa» — JAX è, di fatto, la direzione interna di Google. Il secondo, meno noto: una TPU di generazione corrente costa 2,70 dollari per chip all'ora contro circa undici dollari per GPU su una macchina equivalente dello stesso fornitore. Per chi addestra su larga scala, quella differenza è il vero argomento. E ci sono due limiti che nessuno racconta e che troverete qui: non esiste un percorso ufficiale verso ONNX, e il supporto alle GPU Apple è fermo all'ottobre 2024.

INDICE DEI CONTENUTI

1. Che cos'è JAX, e che cosa non è
   1. La definizione, e perché è importante prenderla alla lettera
   2. Versioni, la numerazione 0.x e una policy di compatibilità insolitamente chiara
   3. Chi lo sviluppa, e la domanda sulla governance
   4. L'adozione, con la metodologia dichiarata
2. Le caratteristiche principali
   1. Le trasformazioni componibili, che sono tutto
   2. Il parallelismo, e una migrazione in corso da conoscere
   3. I vincoli, che sono la parte che si sottovaluta
   4. La compilazione: XLA, StableHLO e la cache
3. I vantaggi competitivi
   1. Il calcolo scientifico differenziabile: qui JAX non ha rivali
   2. La scala, e il conto che la giustifica
   3. Un solo modello mentale, dal notebook al cluster
   4. Keras 3, il ponte per chi non vuole scrivere tutto a mano
   5. Dove il vantaggio si rovescia
4. Costi
   1. La licenza
   2. Le TPU, che sono il motivo economico per cui JAX esiste
   3. Le GPU, per completezza
   4. L'accesso gratuito, con le dovute precisazioni
   5. Il costo vero, che su JAX è più alto che altrove
   6. Tre ordini di grandezza
5. Funzionalità avanzate e integrazioni
   1. Lo stack: che cosa serve oltre a JAX
   2. Pallas: scrivere kernel a mano, e la seconda migrazione in corso
   3. Mettere in produzione: il capitolo più debole
   4. Interoperabilità con PyTorch e NumPy
   5. Hardware: la tabella che decide l'ambiente di lavoro
6. Alternative a JAX e confronto
   1. Il confronto in una tabella
   2. JAX contro PyTorch: la domanda vera
   3. JAX contro TensorFlow, con la citazione che chiude il discorso
   4. E se il problema fosse tabellare?
   5. Come si decide: cinque domande
7. Strategie di crescita usando JAX
   1. Quando ha senso adottarlo, in concreto
   2. Il percorso di adozione che funziona
   3. Per un professionista: che cosa vale questa competenza
   4. Formazione: buona documentazione, nessuna certificazione
   5. Il mercato italiano
8. Tendenze future
   1. Non esiste una roadmap pubblica: la direzione si legge nel changelog
   2. L'hardware, che è il motore di tutta la storia
   3. Agenti e apprendimento per rinforzo
   4. La convergenza sotto il cofano
   5. Che cosa non cambierà
9. Domande frequenti
   1. Che cos'è JAX, e perché non è un framework di deep learning?
   2. Meglio JAX o PyTorch?
   3. JAX è difficile da imparare?
   4. Conviene addestrare su TPU invece che su GPU?
   5. Come si mette in produzione un modello JAX?
   6. Posso usare JAX sul Mac con accelerazione GPU?
   7. Chi usa davvero JAX?
   8. Esiste una certificazione ufficiale JAX?
10. In sintesi

Che cos'è JAX, e che cosa non è

La definizione, e perché è importante prenderla alla lettera

La documentazione ufficiale dichiara tre pilastri: un'API familiare in stile NumPy, trasformazioni componibili di programmi, ed esecuzione su CPU, GPU e TPU con lo stesso codice. La documentazione di Google Cloud aggiunge la frase che spiega la filosofia meglio di qualunque altra: «invece di essere un framework di machine learning monolitico, JAX è deliberatamente ristretto e si concentra su operazioni efficienti su array e trasformazioni di programmi».

Tradotto in pratica: si scrive una funzione matematica in Python usando un'API quasi identica a NumPy, e poi la si decora. Una decorazione la compila, un'altra ne calcola il gradiente, un'altra la applica in parallelo a un intero lotto di dati, un'altra la distribuisce su centinaia di acceleratori. E — questo è il punto — si compongono: il gradiente di una funzione vettorizzata e compilata è ancora una funzione che si può vettorizzare, compilare e distribuire.

Versioni, la numerazione 0.x e una policy di compatibilità insolitamente chiara

La versione stabile corrente è la 0.11.1, rilasciata il 17 agosto 2026, con un ritmo di rilascio mensile e un cambio di numero intermedio circa ogni tre mesi. La licenza è Apache 2.0 e la versione minima di Python richiesta è la 3.12: chi è fermo alla 3.11 non può installare la linea corrente.

Il numero di versione che comincia per zero non è un vezzo: il progetto usa dichiaratamente un modello in cui, finché si resta sotto la versione 1, l'incremento della seconda cifra porta rotture maggiori e quello della terza rotture minori. La cosa notevole — e rara nel settore — è che la policy è scritta e quantificata: una funzione dichiarata obsoleta può essere rimossa non prima di tre mesi, ogni deprecazione va annunciata nel registro delle modifiche e deve emettere un avviso.

Va però letto anche che cosa non è coperto dalle garanzie, perché è un elenco lungo: tutto ciò che comincia con un trattino basso, l'intero pacchetto binario che contiene il compilatore — considerato privato — i moduli storici in via di dismissione e, soprattutto, l'intero spazio dei nomi sperimentale, dove però vivono cose che molti progetti usano in produzione. C'è infine un'avvertenza che vale la pena conoscere prima di far dipendere un collaudo da un numero: i valori numerici esatti delle operazioni non sono garantiti stabili fra una versione e l'altra.

Chi lo sviluppa, e la domanda sulla governance

JAX è sviluppato da Google, fra Google DeepMind e Google Research. Nel settembre 2024 il repository si è spostato dall'organizzazione di Google a una dedicata, e la discussione ufficiale che annunciava il trasferimento è esplicita sul fatto che si trattava di ragioni di infrastruttura di collaudo e di sicurezza, non di governance: «questo spostamento non cambia nulla del modo in cui JAX viene sviluppato».

La conseguenza va detta chiaramente, perché è il rovescio esatto della situazione di PyTorch: non esiste una fondazione indipendente, né una pagina di governance formale. JAX è un progetto di una singola azienda, con licenza permissiva. Per un'azienda che valuta il rischio a cinque anni, è la stessa struttura che aveva TensorFlow prima del marzo 2026 — con la differenza, non piccola, che JAX è oggi la direzione su cui quella stessa azienda sta puntando.

L'adozione, con la metodologia dichiarata

La premessa metodologica vale come nelle altre guide della serie: i download dei pacchetti Python includono integrazione continua, mirror e ricostruzioni di container, quindi non misurano persone. Servono per il confronto relativo. I dati, letti l'11 settembre 2026 sull'ultimo mese:

Libreria	Download nell'ultimo mese	Stelle su GitHub
scikit-learn	214,2 milioni	67,2 mila
PyTorch	73,6 milioni	102 mila
TensorFlow	14,1 milioni	197 mila
JAX	13,9 milioni	36,3 mila

La lettura onesta è questa: JAX ha ormai raggiunto TensorFlow per volume di installazioni mensili — 13,9 milioni contro 14,1 — ma resta a circa un quinto di PyTorch e a un quindicesimo di scikit-learn. Il totale storico è di circa 404 milioni di download, contro i miliardi degli altri tre: è una libreria giovane in crescita, non una libreria di massa. Chi la sceglie deve mettere nel conto una comunità più piccola e meno risposte pronte quando qualcosa non funziona.

Le caratteristiche principali

Le trasformazioni componibili, che sono tutto

È il cuore, e vale la pena elencarle per quello che fanno.

• Compilazione. Marcando una funzione, JAX la traccia una volta e la compila in codice ottimizzato per l'acceleratore, fondendo le operazioni ed eliminando l'overhead dell'interprete Python. È la trasformazione che fa la differenza di prestazioni.

• Differenziazione automatica. Non solo il gradiente: la libreria espone modalità in avanti e all'indietro, matrici jacobiane, hessiane, prodotti vettore-jacobiana in entrambi i versi, e la possibilità di definire derivate personalizzate. È molto più di quanto serva per addestrare una rete, e non è un caso.

• Vettorizzazione automatica. Si scrive la funzione per un esempio, e una trasformazione la applica a un intero lotto senza scrivere cicli né aggiungere dimensioni a mano. È la caratteristica che più spesso fa innamorare chi arriva da NumPy, perché elimina un'intera categoria di codice illeggibile.

• Parallelizzazione. Le funzioni che distribuiscono il calcolo su più dispositivi, di cui parliamo fra poco.

• Ricalcolo per risparmiare memoria. Una trasformazione che, invece di conservare i risultati intermedi per il passaggio all'indietro, li ricalcola: si scambia tempo con memoria, ed è ciò che permette di addestrare modelli che altrimenti non ci starebbero.

La proprietà che conta è la componibilità: sono operatori che prendono una funzione e restituiscono una funzione, quindi si annidano in qualunque ordine. Da qui discende il vantaggio distintivo che vedremo nella sezione sui vantaggi competitivi.

Il parallelismo, e una migrazione in corso da conoscere

Il modello attuale ruota attorno a tre concetti. Ogni array porta con sé un schema di distribuzione che descrive come è ripartito sulle memorie dei diversi dispositivi; i dispositivi sono organizzati in una griglia con assi nominati; e lo schema si esprime dicendo quale asse dell'array vive su quale asse della griglia. La documentazione segnala un dettaglio che ha effetti misurabili: poiché la comunicazione avviene lungo gli assi della griglia, la forma della griglia e l'ordine dei dispositivi determinano le prestazioni di comunicazione.

Su questa base ci sono tre modalità componibili: automatica, in cui si scrive il programma come se ci fosse un solo grande dispositivo e il compilatore decide come ripartirlo e dove inserire le comunicazioni; esplicita, in cui la distribuzione è visibile nei tipi e ispezionabile; e manuale, in cui si scrive il codice dal punto di vista del singolo dispositivo e si inseriscono a mano le operazioni collettive.

Qui c'è la prima delle due migrazioni in corso nel 2026, e riguarda chiunque abbia codice JAX scritto qualche anno fa. La vecchia funzione di parallelizzazione, non è formalmente rimossa ma la sua documentazione dice testualmente di usare l'alternativa moderna, motivandolo con una derivazione automatica più efficiente e una migliore componibilità; con la versione 0.8.0 è passata in manutenzione ed è stata reimplementata sopra le nuove API, e nella linea 0.10 è stata rimossa la sua infrastruttura nativa insieme a un paio di funzioni accessorie. È legacy di fatto, e il codice che la usa va migrato prima che la decisione la prenda qualcun altro.

I vincoli, che sono la parte che si sottovaluta

JAX chiede di scrivere in un certo modo, e non è negoziabile. La documentazione ufficiale dedica un intero documento agli «scogli comuni», il che è già di per sé un'informazione sul salto concettuale richiesto.

• Le funzioni devono essere pure. Tutti i dati entrano dai parametri, tutti i risultati escono dal valore di ritorno. Stampe, variabili globali e stato mutabile producono risultati inattesi o congelati alla prima esecuzione.

• Gli array sono immutabili. Assegnare a un elemento solleva un errore: si usa una sintassi che restituisce una copia modificata.

• Le forme devono essere statiche. Dentro una funzione compilata, la dimensione di un ciclo o il ramo di una condizione non possono dipendere dai valori degli array, solo dalle loro forme. Filtrare un array in base a una condizione — l'operazione più naturale del mondo in NumPy — non si può fare, e si sostituisce con maschere.

• La casualità è esplicita. Niente stato globale del generatore: si crea una chiave e la si divide ogni volta che serve. È più verboso e rende gli esperimenti perfettamente riproducibili.

• La precisione predefinita è a 32 bit, e la doppia precisione va abilitata esplicitamente all'avvio. Per chi fa calcolo scientifico è la prima cosa da configurare.

• Ogni combinazione nuova di forme e tipi causa una ricompilazione. Su codice scritto senza attenzione, il tempo si consuma tutto lì.

La compilazione: XLA, StableHLO e la cache

Sotto le trasformazioni c'è il compilatore. Il percorso documentato è in quattro stadi: la funzione viene tracciata in una rappresentazione intermedia, abbassata al formato StableHLO — lo stesso formato comune di cui parliamo nelle altre guide della serie, prodotto anche da PyTorch e TensorFlow — compilata in un eseguibile e infine eseguita. Ogni stadio è ispezionabile, con l'avvertenza ufficiale che quegli strumenti sono un ausilio al debug, non un'API su cui programmare.

Due cose pratiche che fanno la differenza sui tempi. La prima è la cache di compilazione persistente, che nel 2026 è disattivata per impostazione predefinita e va abilitata indicando una cartella: su progetti con molte esecuzioni ripetute è il primo intervento da fare. Ha però tre limiti dichiarati che vanno conosciuti — la cache è considerata fidata, quindi non va condivisa con utenti non fidati perché consentirebbe l'esecuzione di codice arbitrario; è inefficace con alcune forme di partizionamento personalizzato; e in configurazioni multi-nodo scrive un solo processo, quindi serve uno spazio condiviso.

La seconda è la donazione dei buffer: si dichiara che certi argomenti possono essere sovrascritti dal calcolo, invece di allocare nuova memoria. Su modelli grandi è la differenza fra entrare nella memoria dell'acceleratore e non entrarci.

I vantaggi competitivi

Il calcolo scientifico differenziabile: qui JAX non ha rivali

È il vantaggio vero, e quasi mai viene raccontato per primo perché non riguarda le reti neurali. La combinazione di derivazione automatica, vettorizzazione e compilazione applicata a codice numerico arbitrario — non a un grafo di strati, a codice — permette di fare cose che negli altri framework sono innaturali: derivare un simulatore fisico, ottimizzare i parametri di un modello a equazioni differenziali, calcolare una matrice hessiana completa, fare inferenza bayesiana con campionamento hamiltoniano su un modello scritto a mano.

Non a caso l'ecosistema scientifico costruito su JAX è il più ricco del suo settore: risolutori di equazioni differenziali ordinarie e stocastiche, ottimizzazione non lineare e solutori lineari, programmazione probabilistica e campionamento, dinamica molecolare, fisica differenziabile per la robotica. Sono tutte librerie attive, con rilasci nel 2026.

Chi lavora in ambiti dove il modello è un'equazione e non un'architettura — finanza quantitativa, ingegneria, chimica computazionale, epidemiologia, geofisica — trova in JAX qualcosa che PyTorch può fare ma non fa altrettanto bene, e che scikit-learn non fa affatto. È anche il segmento in cui, in Italia, la concorrenza professionale è più bassa e i committenti sono università, centri di ricerca e aziende manifatturiere.

La scala, e il conto che la giustifica

Il secondo vantaggio è l'addestramento su larga scala, e ha una componente economica misurabile che vediamo nella sezione sui costi: il costo per chip delle TPU è pubblicamente molto inferiore al costo per GPU su macchine equivalenti dello stesso fornitore, e JAX è lo strumento progettato per quelle macchine.

Sul piano tecnico, l'infrastruttura per il multi-host è nativa e non un'aggiunta: un array può estendersi su tutti i processi, e se ogni processo applica la stessa funzione il risultato è — nella formulazione della documentazione — come «programmare contro un unico grande dispositivo». Su TPU la gerarchia fra collegamenti interni veloci e collegamenti fra gruppi di macchine più lenti viene rilevata automaticamente all'inizializzazione, e la griglia dei dispositivi si può costruire tenendone conto. Sono dettagli che su un cluster da migliaia di chip fanno la differenza fra un addestramento che scala e uno che no.

Un solo modello mentale, dal notebook al cluster

È un vantaggio sottile e reale: lo stesso codice gira su CPU, su una GPU, su otto GPU e su un migliaio di chip TPU, cambiando la descrizione della distribuzione e non la logica. Non ci sono due API — una per la prototipazione e una per la produzione — e non c'è un momento in cui si riscrive tutto per andare in scala. Su progetti di ricerca applicata che partono piccoli e possono diventare grandi, è ciò che evita il rifacimento.

Keras 3, il ponte per chi non vuole scrivere tutto a mano

Va detto perché risolve la principale obiezione pratica a JAX — «non ha un'API di alto livello». Keras 3 gira su JAX come backend, quindi si possono costruire modelli con l'API più semplice in circolazione e farli eseguire dal compilatore di JAX. C'è di più, ed è un dettaglio poco noto: l'API di distribuzione di Keras è disponibile solo sul backend JAX, e permette di attivare il parallelismo sui dati in due righe e quello sul modello con una griglia di dispositivi. Per un'azienda che vuole i vantaggi di scala senza formare il team sul modello funzionale, è la strada più corta.

Dove il vantaggio si rovescia

Cinque limiti seri. Chi ne valuta l'adozione deve pesarli tutti.

Primo: la curva di apprendimento. Purezza funzionale, immutabilità, forme statiche, casualità esplicita. Non è difficoltà gratuita — ognuno di quei vincoli esiste perché il compilatore possa fare il suo lavoro — ma è un salto concettuale che la documentazione ufficiale riconosce dedicandogli un documento intero e una guida apposta per chi arriva da PyTorch.

Secondo: l'ecosistema di modelli pre-addestrati è una frazione di quello di PyTorch. Se il vostro progetto comincia scaricando un modello esistente, quel modello con ogni probabilità non è scritto in JAX.

Terzo: non esiste un percorso ufficiale verso ONNX. È una lacuna reale e va detta senza attenuanti: la documentazione ufficiale non ha alcuna pagina sull'argomento, e l'unica strada è un progetto di terze parti. Per un'azienda che ha standardizzato l'inferenza su quel formato, è potenzialmente un motivo di esclusione.

Quarto: su Mac, in pratica, non c'è accelerazione. Il backend per le GPU Apple è dichiarato sperimentale da Apple stessa, non supporta la doppia precisione né i numeri complessi, e — dato più importante di tutti — il pacchetto è fermo alla versione dell'8 ottobre 2024, mentre JAX è arrivato alla 0.11. Su Apple Silicon, oggi, JAX gira sulla CPU: per l'accelerazione locale l'alternativa reale è un altro strumento.

Quinto: la governance. Nessuna fondazione, nessuna roadmap pubblica, numerazione ancora sotto la versione 1, e deprecazioni che possono chiudersi in tre mesi. Per un progetto di ricerca è irrilevante; per un sistema industriale che deve restare fermo cinque anni è un fattore da mettere a contratto, non da ignorare.

Costi

La licenza

Apache 2.0: gratuito, uso commerciale libero, con la consueta clausola sui brevetti. Come per tutta la serie, il costo non è il software: è l'hardware e sono le persone. Su JAX, però, la voce hardware ha una particolarità che vale la pena esaminare.

Le TPU, che sono il motivo economico per cui JAX esiste

Listino pubblico verificato l'11 settembre 2026, prezzo per chip all'ora:

TPU	Regione	Su richiesta	Impegno 1 anno	Impegno 3 anni
v5e	Stati Uniti	1,20 $	0,84 $	0,54 $
v5e	Paesi Bassi	1,56 $	1,09 $	0,70 $
v6e (Trillium)	Stati Uniti	2,70 $	1,89 $	1,22 $
v6e	Paesi Bassi	2,97 $	2,08 $	1,34 $
v7 (Ironwood)	Stati Uniti	12,00 $	8,40 $	5,40 $
v7 (Ironwood)	Regno Unito	13,20 $	9,24 $	5,94 $

Tre letture utili. La prima, per un lettore europeo: le TPU sono disponibili in Europa — la generazione corrente ad Amsterdam, quella di ultima generazione a Londra — con un sovrapprezzo di circa il 10% rispetto agli Stati Uniti, il che ha conseguenze non banali sulla localizzazione dei dati. La seconda: l'impegno triennale taglia il prezzo di circa il 55% su tutte le generazioni, costantemente. La terza, la più interessante: una TPU di generazione corrente a 2,70 dollari per chip all'ora va confrontata con circa undici dollari per GPU all'ora su una macchina a otto acceleratori dello stesso fornitore.

Le GPU, per completezza

Sullo stesso fornitore cloud, sempre a listino pubblico e per macchina a otto acceleratori: circa 88,5 dollari l'ora per otto acceleratori di generazione precedente — cioè circa 11,1 per unità — e circa 64,4 dollari l'ora per otto di ultima generazione, cioè circa 8,1 per unità. Su fornitori specializzati i prezzi sono molto più bassi: 3,49 dollari l'ora per un acceleratore di fascia alta su un fornitore e 3,99 su un altro, con fatturazione al secondo o al minuto.

Il confronto che serve davvero a decidere è quindi in due passaggi. Contro l'hyperscaler, la TPU vince nettamente: 2,70 contro 11,1 dollari. Contro un fornitore specializzato di GPU, il confronto è alla pari — 2,70 contro 3,49 — e si ribalta a favore della TPU solo con l'impegno pluriennale, che porta il prezzo a 1,22. Come sempre, il numero che conta non è il prezzo orario ma il costo per unità di lavoro utile, che dipende da quanto bene il vostro modello sfrutta l'acceleratore: va misurato, non dedotto.

L'accesso gratuito, con le dovute precisazioni

Esiste ed è reale, ma va raccontato con onestà.

Il programma di ricerca sulle TPU offre accesso gratuito a un gruppo dichiarato di oltre mille dispositivi, con supporto esplicito a JAX, in cambio dell'impegno a condividere i risultati — pubblicazioni sottoposte a revisione, codice aperto, articoli. È un'opportunità concreta per università, dottorandi e ricercatori indipendenti italiani, e sorprendentemente poco sfruttata. La durata dell'accesso e le generazioni specifiche di hardware non sono però dichiarate sulla pagina.

Sull'ambiente di notebook gratuito, invece, una precisazione che manca in quasi tutta la divulgazione: non c'è alcuna garanzia pubblicata. Le domande frequenti ufficiali dicono che i tipi di acceleratore disponibili «variano nel tempo», che i limiti «non vengono pubblicati» e che nella versione gratuita l'accesso alle risorse costose è «fortemente limitato». Si può usare per imparare; non ci si può pianificare sopra.

Il costo vero, che su JAX è più alto che altrove

Va detto con franchezza perché è la voce che fa saltare i progetti: su JAX il costo delle persone è superiore a quello di qualunque altro strumento di questa serie. Le competenze sono più rare, la curva è più ripida, la comunità è più piccola e le risposte pronte sono meno. Un'azienda che adotta JAX deve mettere a budget un periodo di formazione reale e — soprattutto — deve chiedersi che cosa succede quando la persona che ha scritto quel codice non c'è più.

È il motivo per cui la raccomandazione di questa guida è netta: JAX si adotta quando c'è una ragione specifica, non per curiosità tecnologica. Le ragioni buone esistono e sono nelle prossime sezioni.

Tre ordini di grandezza

• Ricerca e prototipazione. CPU o una GPU a noleggio per qualche decina di ore al mese, oppure il programma gratuito se si è in ambito accademico. Poche decine o centinaia di euro: il costo dominante è il tempo delle persone per imparare.

• Addestramento di un modello medio su TPU. Una macchina con alcuni chip di generazione corrente, usata a sprazzi: qui il conto si misura in migliaia di euro, e la leva principale non è il prezzo orario ma quante ore si sprecano in esperimenti mal impostati.

• Addestramento su larga scala. Decine o centinaia di chip, dove diventano determinanti l'impegno pluriennale — che dimezza abbondantemente il conto — l'efficienza della pipeline dei dati e la capacità di riprendere da un punto di controllo dopo un guasto. È il terreno per cui JAX è stato progettato, ed è anche l'unico in cui la sua complessità si ripaga da sola.

Funzionalità avanzate e integrazioni

Lo stack: che cosa serve oltre a JAX

Poiché JAX è deliberatamente ristretto, un progetto reale ne usa almeno quattro pezzi. Versioni verificate all'11 settembre 2026:

Libreria	A che cosa serve	Versione e stato
Flax	Definire reti neurali	0.12.9 (18 ago 2026) — attiva
Optax	Ottimizzatori e trasformazioni del gradiente	0.2.8 (20 mar 2026) — attiva
Orbax	Punti di controllo distribuiti e asincroni	0.12.4 (12 ago 2026) — attiva
Grain	Pipeline dei dati deterministica e ripristinabile	0.2.18 (17 giu 2026) — attiva
Chex	Utilità di test e asserzioni	0.1.92 (12 giu 2026) — attiva
Equinox	Alternativa a Flax, più vicina al Python idiomatico	0.13.8 (5 mag 2026) — attiva
Haiku	La libreria storica di DeepMind	0.0.17 — in manutenzione dal luglio 2023: DeepMind raccomanda Flax per i nuovi progetti
MaxText	Riferimento per addestrare modelli linguistici su TPU e GPU	0.2.4 (21 ago 2026) — attiva, ma il ramo principale è dichiarato non pronto per la produzione

Su Flax serve una precisazione, perché è la domanda più frequente di chi comincia adesso. Esistono due API: quella storica e quella nuova. La storica non è deprecata — la documentazione dice esplicitamente che non lo sarà nel prossimo futuro, perché la maggior parte degli utenti la usa ancora — ma i nuovi progetti sono incoraggiati a usare quella nuova, descritta come il risultato di anni di esperienza in una forma più semplice. Esiste una guida di migrazione.

Pallas: scrivere kernel a mano, e la seconda migrazione in corso

Quando l'ottimizzazione automatica del compilatore non basta — tipicamente per meccanismi di attenzione personalizzati, quantizzazione o sparsità — si scende al livello dei kernel, cioè del codice che gira direttamente sull'acceleratore. JAX ha un'estensione dedicata che permette di scriverli in Python, con controllo esplicito sulla memoria e sul pipelining.

Va presa per quello che è, e la documentazione lo dichiara: è in sviluppo attivo e non stabile, ha un registro delle modifiche separato, e il supporto hardware è selettivo — sulle GPU solo le due generazioni più recenti di fascia datacenter, e sulle TPU le due unità di calcolo principali. I due backend condividono i concetti ma hanno API diverse.

E qui c'è la seconda migrazione del 2026: nella versione 0.11.0, di luglio 2026, il backend basato sulla tecnologia di compilazione di terze parti è stato dichiarato obsoleto, con invito a passare a quello sviluppato internamente; nella 0.11.1 le operazioni già scritte per il vecchio backend sono state marcate per la rimozione. Chi ha kernel scritti con quella tecnologia ha un lavoro davanti.

Mettere in produzione: il capitolo più debole

È giusto dirlo apertamente perché è il punto in cui JAX è più indietro rispetto a PyTorch.

Il percorso nativo serializza la funzione compilata in un artefatto portabile nel formato comune, con garanzie insolitamente precise e insolitamente strette: un artefatto è eseguibile da un runtime fino a sei mesi più nuovo e fino a tre settimane più vecchio di quello che l'ha prodotto. Nel 2026 quella scadenza è diventata effettiva: nella linea 0.11 gli artefatti esportati prima del 15 gennaio 2026 non sono più deserializzabili, salvo una variabile d'ambiente di emergenza. Tradotto: un modello JAX esportato non è un artefatto che si mette in un cassetto per tre anni.

Il percorso verso l'ecosistema TensorFlow — utile per chi ha già infrastruttura di servizio basata su quel formato, e obbligatorio per l'on-device — esiste, è in uso dal 2020 e la sua documentazione dichiara che la grande maggioranza degli utenti non incontra problemi. Ha però una caratteristica che va messa nel capitolato: vive ancora nello spazio dei nomi sperimentale, cioè fuori dalle garanzie di compatibilità. E la documentazione ufficiale per il deployment sul dispositivo passa proprio da lì, definendolo essa stessa «una funzionalità sperimentale».

Verso ONNX non esiste alcun percorso ufficiale. Nessuna pagina nella documentazione, nessun convertitore mantenuto dal progetto: l'unica strada è una libreria di terze parti. Per un'azienda che ha standardizzato l'inferenza su quel formato — scelta molto comune, e spesso saggia — è un fattore di esclusione da valutare prima, non dopo.

Per il servizio di modelli linguistici, la raccomandazione ufficiale di Google Cloud è un motore di inferenza esterno nella sua versione per TPU — lo stesso, non a caso, che nella nostra guida a PyTorch risulta lo standard di fatto. È un buon esempio di come, a valle, i due mondi convergano.

Interoperabilità con PyTorch e NumPy

Con NumPy la conversione è diretta in entrambi i versi. Con PyTorch esiste un modulo ufficiale per lo scambio di array senza copia basato su uno standard comune, che permette di passare tensori fra i due mondi sullo stesso dispositivo; i pacchetti che automatizzano la conversione di intere funzioni sono invece progetti di terze parti, senza supporto ufficiale. È abbastanza per costruire una pipeline mista — per esempio caricamento dati con gli strumenti di PyTorch e calcolo in JAX — non abbastanza per considerare i due framework intercambiabili.

Hardware: la tabella che decide l'ambiente di lavoro

Acceleratore	Linux	macOS	Windows
CPU	sì	sì	sì
GPU NVIDIA	sì	no	no in nativo, sperimentale via sottosistema Linux
TPU	sì	—	—
GPU AMD	sì (solo x86)	—	sperimentale via sottosistema Linux
GPU Apple	—	sperimentale, e di fatto fermo	—
GPU Intel	sperimentale	—	no

Due righe meritano un commento. Windows nativo non supporta le GPU NVIDIA: serve il sottosistema Linux, e anche lì il supporto è dichiarato sperimentale. È un limite che pesa in molte aziende italiane, dove il parco macchine degli sviluppatori è Windows. Le GPU Apple sono formalmente supportate in modo sperimentale, ma il pacchetto non viene aggiornato dall'ottobre 2024: nella pratica, su Mac si sviluppa su CPU. Sono i due vincoli che decidono l'ambiente di lavoro di un team, e vanno verificati prima di cominciare, non dopo.

Alternative a JAX e confronto

Il confronto in una tabella

Strumento	Versione (set. 2026)	Quando conviene rispetto a JAX
JAX	0.11.1 (17 ago 2026)	Calcolo scientifico differenziabile, addestramento su larga scala, TPU, controllo fine della distribuzione
PyTorch	2.14.0 (2 set 2026)	Quasi sempre: ecosistema di modelli, tooling, forme dinamiche, debug immediato, GPU NVIDIA in azienda
TensorFlow	2.21.0 (6 mar 2026)	Solo per basi di codice esistenti e per il deployment su dispositivo con il suo runtime
Keras 3	3.15.1 (29 lug 2026)	API di alto livello sopra JAX: produttività senza il modello funzionale. È anche l'unico backend con l'API di distribuzione
scikit-learn	1.9.1 (10 set 2026)	Dati tabellari: nessuna sovrapposizione reale, e nella maggioranza dei casi aziendali è la risposta giusta
MLX	0.32.2 (25 ago 2026)	Su Apple Silicon, dove il backend Apple di JAX è fermo da quasi due anni
Julia	—	Calcolo scientifico differenziabile senza il livello Python, per chi accetta un linguaggio diverso

JAX contro PyTorch: la domanda vera

Il confronto si decide su tre assi, e per la maggior parte dei lettori il risultato è a favore di PyTorch. Vale la pena dirlo apertamente in una guida su JAX, perché è la cosa utile.

L'ecosistema. PyTorch ha i modelli pre-addestrati, le librerie di adattamento, i motori di servizio, i tutorial e il codice degli articoli scientifici. Se il vostro progetto comincia da un modello che esiste già, comincia su PyTorch.

La produttività immediata. PyTorch permette forme dinamiche, control flow che dipende dai dati, stampe in mezzo alla rete, stato mutabile. JAX no, e ogni «no» è un vincolo che il compilatore usa per andare più veloce ma che il programmatore paga in tempo.

La scala e la matematica. Qui JAX recupera tutto. Se dovete derivare un simulatore, calcolare hessiane, fare inferenza bayesiana, o distribuire un addestramento su migliaia di chip con controllo esplicito su come ogni tensore è ripartito, JAX fa cose che PyTorch fa con più fatica o non fa.

La regola pratica che ne deriva: PyTorch è la scelta predefinita, JAX è la scelta motivata. Se non sapete articolare in una frase perché vi serve JAX, vi serve PyTorch.

JAX contro TensorFlow, con la citazione che chiude il discorso

Il confronto è ormai storico più che tecnico. Il 6 marzo 2026, annunciando la versione 2.21, Google ha dichiarato che si concentrerà esclusivamente su correzioni di sicurezza e di difetti e aggiornamenti delle dipendenze, e ha scritto: «Mentre TensorFlow continua a garantire stabilità per la produzione, raccomandiamo di esplorare i nostri aggiornamenti più recenti su Keras 3, JAX e PyTorch per il nuovo lavoro sull'intelligenza artificiale generativa».

Per chi legge questa guida significa due cose. La prima: JAX è la direzione dichiarata di Google per il calcolo numerico e l'addestramento, ed è ciò che spiega l'investimento sullo stack, sulle TPU e sull'infrastruttura di scala. La seconda, più pratica: chi ha codice TensorFlow e sta valutando dove andare ha due destinazioni ufficialmente indicate, e la scelta fra le due dipende dai criteri della sezione precedente. Il percorso più morbido per chi arriva da lì è quasi sempre Keras 3 con backend JAX: si conserva l'API, si cambia il motore.

E se il problema fosse tabellare?

Va ripetuto anche qui, perché è il promemoria che fa risparmiare più soldi in assoluto: se i dati sono righe e colonne, non serve niente di tutto questo. Un gradient boosting su CPU si addestra in minuti, si spiega a chi decide e costa una frazione. JAX e i dati tabellari non si sovrappongono: sono mondi diversi, e confonderli è l'errore più costoso che un responsabile tecnico possa fare in questo campo.

Come si decide: cinque domande

1. Il vostro modello è un'equazione o un'architettura? Se è un'equazione — un simulatore, un sistema dinamico, un modello probabilistico — JAX è probabilmente lo strumento giusto.

2. Addestrate su TPU o su molte centinaia di acceleratori? Se sì, JAX è progettato per quello.

3. Partite da un modello pre-addestrato esistente? Se sì, PyTorch.

4. La vostra inferenza è standardizzata su un formato portabile? Se quel formato è ONNX, verificate subito: non c'è un percorso ufficiale.

5. Chi manterrà questo codice fra due anni, e ha già scritto codice funzionale? Su JAX è la domanda che pesa più di tutte le altre messe insieme.

Strategie di crescita usando JAX

Quando ha senso adottarlo, in concreto

Quattro situazioni in cui la scelta si difende da sola davanti a un comitato.

Uno: il modello è fisico o matematico. Simulazioni differenziabili, calibrazione di modelli a equazioni differenziali, ottimizzazione di progetto, inferenza bayesiana. Qui la combinazione di derivazione e vettorizzazione su codice arbitrario non ha equivalenti pratici.

Due: si addestra su larga scala e il conto dell'hardware è una voce di bilancio. Con l'impegno pluriennale il prezzo per chip scende del 55%, e il controllo esplicito sulla distribuzione permette di sfruttare l'hardware meglio di quanto faccia un'astrazione automatica.

Tre: si sta migrando da TensorFlow. È la destinazione indicata dal produttore, e con Keras 3 il passaggio conserva l'API di alto livello.

Quattro: si lavora in ricerca e si può accedere al programma gratuito sulle TPU. Per un'università o un centro di ricerca italiano, avere accesso a hardware di quella classe senza costi, in cambio della pubblicazione dei risultati, è un'opportunità che vale più di molti bandi.

Il percorso di adozione che funziona

1. Cominciate dalla parte matematica, non dalla rete. Prendete un calcolo che già fate in NumPy e trasformatelo: compilate, vettorizzate, derivate. Si capisce il valore in un'ora.

2. Imparate i vincoli prima delle librerie. Purezza, immutabilità, forme statiche, chiavi di casualità. La documentazione ufficiale ha un documento dedicato agli scogli comuni: è la lettura con il miglior rapporto fra tempo e dolore risparmiato.

3. Usate Keras 3 se il team non è pronto per il modello funzionale: si ottiene il motore senza la curva.

4. Fissate le versioni di tutto, componente per componente, e non affidatevi al meta-pacchetto finché non torna allineato.

5. Decidete presto come uscirà il modello. Formato nativo con la sua scadenza, percorso verso l'altro ecosistema per l'on-device, o motore di inferenza esterno. È il capitolo più debole dello stack e va progettato all'inizio.

6. Mettete a budget la formazione, e più di una persona. Un progetto JAX con un solo esperto è un progetto a termine.

Per un professionista: che cosa vale questa competenza

È la competenza più rara della serie e quella con la nicchia più difendibile, a patto di sceglierla con cognizione. Non ha senso impararla per fare quello che fanno tutti: ha senso per due profili.

Il primo è chi lavora su problemi scientifici e ingegneristici — ottimizzazione, simulazione, modelli probabilistici, finanza quantitativa — dove la capacità di derivare e vettorizzare codice numerico arbitrario è un vantaggio tecnico difficile da replicare, e dove i committenti sono centri di ricerca e industria manifatturiera più che startup.

Il secondo è chi lavora su addestramento su larga scala: distribuzione esplicita, punti di controllo asincroni, pipeline dei dati deterministiche, kernel personalizzati. Sono competenze richieste da poche aziende, che però pagano bene perché trovano pochissimi candidati.

Le competenze trasferibili restano quelle di sempre e valgono anche qui: dati, valutazione, messa in produzione, monitoraggio. Ma su JAX c'è una competenza specifica che vale la pena nominare, perché nessuno la insegna: saper diagnosticare perché una funzione ricompila. È il problema numero uno dei progetti reali, e chi lo sa risolvere fa risparmiare giornate di calcolo.

Formazione: buona documentazione, nessuna certificazione

La documentazione ufficiale è organizzata in un percorso a livelli che va dai concetti di base fino all'architettura interna, passando per prestazioni e scalabilità, derivazione avanzata, kernel personalizzati e sistemi distribuiti. È sequenziale, è di buona qualità ed è gratuita; accanto ci sono il documento sugli scogli comuni e una guida dedicata a chi arriva da PyTorch.

Due difetti oggettivi, riscontrati durante la verifica per questo articolo. Il primo: la documentazione viene riorganizzata spesso e diversi indirizzi storici sono ormai redirezioni, il che rompe i collegamenti esterni e le vecchie risposte trovate online. Il secondo, già detto: il meta-pacchetto promosso dalla documentazione è fermo a una versione di dieci mesi fa.

E una differenza rispetto al resto della serie che vale la pena segnalare: non esiste alcuna certificazione ufficiale JAX. PyTorch ne ha lanciata una nel giugno 2026, scikit-learn ne ha una dal 2024, TensorFlow ne aveva una ora chiusa. Su JAX non c'è, e non risulta annunciata.

Il mercato italiano

Come nelle altre guide della serie: non esistono dati pubblici con metodologia dichiarata sulla domanda di competenze JAX in Italia, né sulle retribuzioni per singola libreria. Chi li cita usa aggregatori di annunci senza campione né metodo.

Il quadro generale — mercato italiano dell'intelligenza artificiale a 1,8 miliardi di euro nel 2025 in crescita del 50%, domanda di competenze negli annunci in crescita del 93%, secondo l'Osservatorio del Politecnico di Milano — va letto qui con una precisazione onesta: la quota di quel mercato che richiede JAX è piccola. Il 54% è machine learning classico, dove serve scikit-learn; buona parte del resto parte da modelli pre-addestrati, dove serve PyTorch.

Ma proprio per questo la nicchia è interessante: pochi progetti, pochissimi professionisti, committenti con budget — centri di ricerca, aziende manifatturiere con problemi di simulazione, gruppi che addestrano modelli propri. È una scelta di posizionamento, non una scelta di volume, e va fatta sapendolo.

Tendenze future

Non esiste una roadmap pubblica: la direzione si legge nel changelog

È un fatto, ed è anche una differenza di stile rispetto agli altri progetti della serie: JAX non pubblica una roadmap. Non c'è una pagina di programma, non c'è un documento di indirizzo. Per capire dove sta andando bisogna leggere il registro delle modifiche — cosa che abbiamo fatto per questa guida — e ciò che emerge sono cinque direttrici chiare.

• Consolidamento sul modello di distribuzione moderno, con lo smantellamento progressivo della vecchia funzione di parallelizzazione: passata in manutenzione con la 0.8 e privata della sua infrastruttura nativa nella linea 0.10.

• Consolidamento dei kernel su un unico backend interno, con la deprecazione di quello basato su tecnologia di terze parti nella 0.11.0 di luglio 2026.

• Irrigidimento delle garanzie di esportazione: gli artefatti hanno ora una scadenza effettiva, e quelli anteriori al 15 gennaio 2026 non si caricano più.

• Sicurezza del calcolo distribuito, con l'introduzione della cifratura reciproca nel servizio di coordinamento fra processi nella 0.11.1.

• Innalzamento della base Python alla 3.12 e supporto stabile all'interprete senza blocco globale — un dettaglio tecnico che conta molto per chi fa calcolo parallelo in Python.

Due di queste direttrici sono migrazioni obbligate per chi ha già codice, e vale la pena averlo in chiaro: chi usa la vecchia parallelizzazione e chi ha scritto kernel con il vecchio backend ha un lavoro da pianificare entro i prossimi cicli di rilascio.

L'hardware, che è il motore di tutta la storia

La generazione corrente di TPU è stata annunciata nell'aprile 2025 e dichiarata la prima progettata specificamente per l'inferenza, con configurazioni fino a 9.216 chip raffreddati a liquido. Con l'annuncio di disponibilità generale del novembre 2025 Google ha dichiarato dieci volte le prestazioni di punta rispetto alla generazione di due passi precedente e quattro volte per chip rispetto alla precedente, oltre a una quantità di memoria condivisa nell'ordine del petabyte per configurazione.

Nello stesso annuncio compaiono due elementi che riguardano direttamente chi legge: il miglioramento degli strumenti di addestramento e post-addestramento dello stack, e il fatto che il motore di inferenza più diffuso permette ormai di «passare da GPU a TPU con modifiche minime di configurazione». Quest'ultimo punto è strategicamente importante, perché riduce il costo di provare: non serve riscrivere per valutare se la TPU conviene.

Una precisazione di metodo: circolano notizie su una generazione successiva di TPU annunciata nella primavera 2026. Non abbiamo trovato conferma su fonti ufficiali e quindi questa guida non la riporta.

Agenti e apprendimento per rinforzo

Il progetto JAX in sé non ha una posizione dichiarata su agenti e modelli linguistici — coerentemente con la sua natura di libreria numerica ristretta. Lo stack costruito attorno, però, sì: esiste una libreria ufficiale per il post-addestramento e l'allineamento che nel dicembre 2025 ha rilasciato funzionalità dedicate all'apprendimento per rinforzo in contesti agentici e supporta gli algoritmi di riferimento del settore.

La lettura per chi decide: se il vostro interesse è costruire agenti, JAX non è il punto di ingresso — lo è PyTorch, con il suo ecosistema. Se invece il vostro interesse è addestrare o riallineare modelli propri su hardware Google, allora questo stack è esattamente il posto dove quel lavoro si sta facendo.

La convergenza sotto il cofano

C'è una tendenza di fondo che attraversa tutta questa serie di guide e che qui si vede meglio che altrove. JAX, PyTorch e TensorFlow producono tutti e tre lo stesso formato intermedio, che i compilatori consumano per generare codice per CPU, GPU, TPU e acceleratori vari. Il motore di inferenza più diffuso per i modelli linguistici gira su GPU e su TPU. Il runtime per l'on-device accetta modelli da tutti e tre.

La conseguenza, per un'azienda, è liberatoria: la scelta del framework sta diventando meno irreversibile di quanto sembri. Non è gratuita — il codice va riscritto — ma il livello sotto è sempre più comune, e questo riduce il rischio strategico di ogni singola decisione. La domanda giusta da porre a un fornitore, come scrivevamo nella guida a PyTorch, non è più «quale framework usate» ma «in quale formato mi consegnate il modello e con quale runtime lo servo».

Che cosa non cambierà

La qualità dei dati, la correttezza della validazione, il monitoraggio in esercizio, il costo dell'inferenza e la conformità normativa: i soliti cinque, che valgono su qualunque strumento. Su JAX se ne aggiunge un sesto, specifico e spesso sottovalutato: il costo di avere in squadra persone che sanno usarlo. È la voce che determina se l'adozione è un vantaggio competitivo o un debito tecnico con una sola persona che lo regge.

Domande frequenti

Che cos'è JAX, e perché non è un framework di deep learning?

Perché non contiene reti neurali. La documentazione ufficiale lo definisce «una libreria Python per il calcolo su array orientato agli acceleratori e per la trasformazione di programmi»: fornisce un'API quasi identica a NumPy e un gruppo di trasformazioni componibili — compilare, derivare, vettorizzare, distribuire, ricalcolare per risparmiare memoria — che si applicano a normale codice numerico Python. Strati, ottimizzatori, punti di controllo e pipeline dei dati stanno in librerie separate. È una scelta dichiarata: la documentazione di Google Cloud scrive che «invece di essere un framework monolitico, JAX è deliberatamente ristretto».

Meglio JAX o PyTorch?

Per la maggior parte dei progetti, PyTorch, e conviene dirlo anche in una guida su JAX: ha l'ecosistema dei modelli pre-addestrati, il tooling, le forme dinamiche e il debug immediato. JAX vince su due terreni precisi: il calcolo scientifico differenziabile — derivare un simulatore, calcolare hessiane, fare inferenza bayesiana su codice scritto a mano — e l'addestramento su larga scala con controllo esplicito di come ogni tensore è distribuito, in particolare su TPU. La regola pratica: PyTorch è la scelta predefinita, JAX è la scelta motivata. Se non sapete dire in una frase perché vi serve JAX, vi serve PyTorch.

JAX è difficile da imparare?

Ha una curva più ripida, e non per difficoltà gratuita: impone funzioni pure, array immutabili, forme statiche e generazione di numeri casuali esplicita, perché sono i vincoli che permettono al compilatore di lavorare. Chi arriva da NumPy o da PyTorch inciampa sempre sulle stesse tre cose: l'assegnazione a un elemento, che qui solleva un errore; il generatore casuale globale, che qui non esiste; e il filtraggio di un array in base a una condizione, che dentro una funzione compilata non si può fare. C'è poi un tranello che non fa rumore ed è il più insidioso: l'indicizzazione fuori dai limiti non genera errore, restituisce l'ultimo elemento. La documentazione ufficiale dedica un intero documento agli «scogli comuni» e una guida apposita a chi arriva da PyTorch.

Conviene addestrare su TPU invece che su GPU?

Sui prezzi di listino verificati a settembre 2026, la risposta dipende con chi confrontate. Una TPU di generazione corrente costa 2,70 dollari per chip all'ora su richiesta, contro circa 11 dollari per GPU all'ora su una macchina a otto acceleratori dello stesso fornitore cloud: lì la TPU vince nettamente. Contro un fornitore specializzato di GPU, che parte da circa 3,49 dollari l'ora, il confronto è alla pari — e si sposta a favore della TPU con l'impegno pluriennale, che taglia il prezzo di circa il 55% portandolo a 1,22. Attenzione però: il numero che conta è il costo per unità di lavoro utile, che dipende da quanto il vostro modello sfrutta l'acceleratore, e va misurato. Esiste inoltre un programma di ricerca che offre accesso gratuito alle TPU in cambio della pubblicazione dei risultati.

Come si mette in produzione un modello JAX?

È il capitolo più debole dello stack e va progettato all'inizio. Il percorso nativo serializza la funzione compilata in un formato portabile con garanzie precise e strette: un artefatto è eseguibile da un runtime fino a sei mesi più nuovo e tre settimane più vecchio, e nel 2026 quella scadenza è diventata effettiva — gli export anteriori al 15 gennaio 2026 non si caricano più. Per l'ecosistema TensorFlow e per il deployment sul dispositivo esiste un convertitore in uso dal 2020, che però vive ancora nello spazio dei nomi sperimentale. Verso ONNX non esiste alcun percorso ufficiale: è una lacuna reale, e per chi ha standardizzato l'inferenza su quel formato può essere un motivo di esclusione. Per i modelli linguistici, la raccomandazione ufficiale è un motore di inferenza esterno nella sua versione per TPU.

Posso usare JAX sul Mac con accelerazione GPU?

In pratica no. Il backend per le GPU Apple esiste ed è dichiarato sperimentale da Apple stessa — non supporta la doppia precisione né i numeri complessi — ma il dato decisivo è un altro: il pacchetto è fermo alla versione dell'8 ottobre 2024, mentre JAX è arrivato alla 0.11.1 nell'agosto 2026. Su Apple Silicon, oggi, JAX gira sulla CPU. Per l'accelerazione locale su Mac l'alternativa reale è MLX. Da sapere anche che Windows nativo non supporta le GPU NVIDIA con JAX: serve il sottosistema Linux, e anche lì il supporto è dichiarato sperimentale.

Chi usa davvero JAX?

Limitandosi alle fonti ufficiali: la scheda modello di Gemma 3 dichiara esplicitamente che l'addestramento è stato fatto con JAX; il repository di AlphaFold 3 lo elenca fra le dipendenze; e un post ufficiale del Google Developers Blog del novembre 2025 nomina Anthropic, xAI e Apple fra chi costruisce su questo stack, portando anche casi con numeri dichiarati. Su Apple c'è un riscontro indipendente, perché il suo framework interno di addestramento è costruito su JAX. Una precisazione che quasi tutti gli articoli sbagliano: non esiste una dichiarazione ufficiale di Google che affermi che Gemini sia addestrato con JAX — è documentato l'uso dell'infrastruttura di orchestrazione, non del framework.

Esiste una certificazione ufficiale JAX?

No, e non risulta annunciata. È una differenza rispetto al resto della serie: PyTorch ha lanciato una certificazione ufficiale nel giugno 2026, scikit-learn ne ha una dal 2024 su tre livelli, e TensorFlow ne aveva una, oggi chiusa. Per JAX ci sono la documentazione ufficiale — organizzata in un percorso a livelli dai concetti di base all'architettura interna, gratuita e di buona qualità — il documento sugli scogli comuni e una guida dedicata a chi arriva da PyTorch. Due difetti oggettivi da conoscere: la documentazione viene riorganizzata spesso, con molti indirizzi storici ormai redirezionati, e il meta-pacchetto ufficiale che dovrebbe installare l'intero stack è fermo a ottobre 2025, mentre JAX è alla 0.11: conviene installare i componenti singolarmente.

In sintesi

JAX è la scelta giusta in due casi, e sono casi precisi. Quando il modello è matematica — un simulatore da derivare, un sistema di equazioni da calibrare, un modello probabilistico da campionare — perché la combinazione di derivazione automatica e vettorizzazione su codice arbitrario non ha equivalenti pratici. E quando si addestra su larga scala, in particolare su TPU, dove il costo per chip è pubblicamente competitivo e il controllo esplicito sulla distribuzione permette di sfruttare l'hardware meglio di un'astrazione automatica. Fuori da questi due casi, per quasi tutti, la risposta è PyTorch — o, se i dati sono tabellari, scikit-learn.

Le tre cose da portare via, se ne resta spazio solo per tre. Primo: JAX non è un framework di deep learning, è un compilatore per programmi numerici, e capirlo cambia completamente il criterio con cui lo si valuta. Secondo: i due limiti che nessuno racconta sono reali — non c'è un percorso ONNX ufficiale e su Mac l'accelerazione è ferma dall'ottobre 2024 — e vanno verificati prima di adottarlo, non dopo. Terzo: il costo dominante non è l'hardware, sono le persone, perché su questo strumento la curva è ripida e le competenze sono rare. Il che è, insieme, il rischio principale dell'adozione e la ragione per cui chi le ha trova poca concorrenza.

Se stai valutando come portare il machine learning dentro la tua azienda, o hai un progetto di simulazione, ottimizzazione o addestramento su larga scala da impostare, su FreelanceWWW trovi professionisti che lavorano su dati, modelli e messa in produzione: pubblica il tuo progetto e ricevi le offerte dai freelance che quel lavoro lo hanno già fatto. Questa guida chiude una serie di quattro: le altre sono su PyTorch, TensorFlow e scikit-learn; sul profilo professionale trovi il freelance Python per l'intelligenza artificiale.