This library is under active development and the API is subject to change.
| Section | Description |
|---|---|
| Introduction | What is AXLearn? |
| Getting Started | Getting up and running with AXLearn. |
| Concepts | Core concepts and design principles. |
| CLI User Guide | How to use the CLI. |
| Infrastructure | Core infrastructure components. |
Welcome to AXLearn, a powerful deep learning library built on top of JAX and XLA. AXLearn is designed to support the development of large-scale deep learning models and offers a unique approach to tackling the software engineering challenges that arise during model development and maintenance. AXLearn is a library built on top of JAX and XLA to support development of large-scale deep learning models.
AXLearn takes an object-oriented approach to the software engineering challenges that arise from building, iterating, and maintaining models. The configuration system of the library lets users compose models from reusable building blocks and integrate with other libraries such as Flax and Hugging Face transformers.
-
AXLearn is built to scale.
-
AXLearn is a powerful machine learning tool that can train models with hundreds of billions of parameters across thousands of accelerators, ensuring high utilization. It offers efficient deployment and management of jobs and data, utilizing GSPMD and adopting a global computation paradigm, allowing computation to be described on a virtual global computer.
-
It is also designed to run on public clouds and provides tools to deploy and manage jobs and data. Built on top of GSPMD, AXLearn adopts a global computation paradigm to allow users to describe computation on a virtual global computer rather than on a per-accelerator basis.
-
-
Versatile Applications.
- AXLearn isn't limited to a specific domain. It supports a wide range of applications, including natural language processing, computer vision, and speech recognition. Additionally, it comes with baseline configurations for training state-of-the-art models. For in-depth information on AXLearn's core components and design, please refer to the Concepts section. If you're eager to dive right in, check out Getting Started.
Ready to harness the power of AXLearn? The Getting Started guide will walk you through the process of setting up AXLearn and embarking on your deep learning journey.
To truly grasp the core components and design principles of AXLearn, head over to the Concepts section, where you'll find valuable insights into our library's inner workings.
Learn how to make the most of AXLearn's Command Line Interface (CLI) in the CLI User Guide.
Understand the key infrastructure components of AXLearn by exploring the Infrastructure section. This is where you'll discover how AXLearn handles the heavy lifting behind the scenes.