AXLearn es una biblioteca de aprendizaje profundo creada por Apple basada en JAX y XLA. Emplea un enfoque de ingeniería de software orientado a objetos para abordar los desafíos que plantea el desarrollo de modelos de aprendizaje profundo a gran escala. Su sistema de configuración permite a los usuarios componer modelos a partir de bloques de construcción reutilizables e integrarse con otras bibliotecas, como Flax y Hugging Face Transformers. AXLearn está diseñado para escalar el entrenamiento, permitiendo el entrenamiento eficiente de modelos con decenas de miles de millones de parámetros en miles de aceleradores, siendo ideal para la implementación en la nube pública. También emplea un paradigma de computación global, permitiendo a los usuarios describir los cálculos en una máquina virtual global en lugar de en un solo acelerador. AXLearn admite una amplia gama de aplicaciones, incluyendo el procesamiento del lenguaje natural, la visión artificial y el reconocimiento de voz, e incluye configuraciones de referencia para entrenar modelos de vanguardia.