AXLearn ist eine von Apple entwickelte Deep-Learning-Bibliothek, die auf JAX und XLA basiert und mithilfe einer objektorientierten Herangehensweise die Herausforderungen der Softwareentwicklung bei der Entwicklung großer Deep-Learning-Modelle löst. Ihr Konfigurationssystem ermöglicht es Nutzern, Modelle aus wiederverwendbaren Bausteinen zusammenzusetzen und mit anderen Bibliotheken (wie Flax und Hugging Face Transformers) zu integrieren. AXLearn zielt auf die Skalierung des Trainings ab und unterstützt das effiziente Training von Modellen mit mehreren Milliarden Parametern auf Tausenden von Beschleunigern, ideal für den Einsatz in der Public Cloud. Es verwendet ein globales Berechnungsmodell, das es Nutzern erlaubt, Berechnungen auf einem globalen virtuellen Computer zu beschreiben, anstatt auf einzelnen Beschleunigern. AXLearn unterstützt ein breites Spektrum an Anwendungen, darunter die Verarbeitung natürlicher Sprache, Computer Vision und Spracherkennung, und enthält Baseline-Konfigurationen zum Trainieren von hochmodernen Modellen.