axlearn

Why we included this project

AXLearn is Apple's deep learning library built on JAX and XLA. Its design is object-oriented and configuration-driven: models get assembled from reusable building blocks that can plug into Flax or Hugging Face transformers, so it works alongside a lot of existing tooling. What separates it from other JAX frameworks is how it handles scale. Instead of reasoning about individual accelerators, it treats the whole cluster as one virtual machine, which lets it train models with hundreds of billions of parameters across thousands of devices at high utilization. It also ships baseline configurations for NLP, vision, and speech, along with tooling for running jobs and managing data on public clouds. Teams that already build on JAX and expect to move from small experiments to very large distributed training runs will find this worth a close look.