AXLearnは、AppleがJAXとXLAを基盤として構築した深層学習ライブラリです。オブジェクト指向のアプローチを採用することで、大規模な深層学習モデル開発におけるソフトウェアエンジニアリング上の課題を解決します。その構成システムにより、ユーザーは再利用可能なビルディングブロックからモデルを組み合わせ、FlaxやHugging Face Transformersなどの他のライブラリと統合することができます。AXLearnはトレーニング規模の拡張を目指しており、数百億パラメータのモデルを数千個のアクセラレータ上で効率的にトレーニングすることを可能にし、パブリッククラウド上での展開に適しています。また、グローバル計算パラダイムを採用することで、ユーザーは個々のアクセラレータではなく、グローバル仮想マシン上の計算を記述できます。AXLearnは自然言語処理、コンピュータビジョン、音声認識など幅広い用途に対応し、最先端モデルのトレーニングのためのベースライン構成を含んでいます。