AXLearn é uma biblioteca de aprendizado profundo desenvolvida pela Apple, baseada em JAX e XLA. Ela utiliza uma abordagem orientada a objetos para solucionar os desafios de engenharia de software no desenvolvimento de modelos de aprendizado profundo em larga escala. Seu sistema de configuração permite que os usuários combinem modelos a partir de blocos de construção reutilizáveis e integrem-no com outras bibliotecas (como Flax e Hugging Face Transformers). O AXLearn foi projetado para escalar o treinamento, suportando o treinamento eficiente de modelos com dezenas de bilhões de parâmetros em milhares de aceleradores, sendo ideal para implantação em nuvem pública. Ele também emprega um paradigma de computação global, permitindo que os usuários descrevam cálculos em uma máquina virtual global, em vez de em aceleradores individuais. O AXLearn suporta uma ampla gama de aplicações, incluindo processamento de linguagem natural, visão computacional e reconhecimento de voz, e inclui configurações base para o treinamento de modelos de última geração.