google/paxml

Pax は、大規模モデルのトレーニングのための Jax ベースの機械学習フレームワークです。高度かつ完全に設定可能な実験と並列処理を可能にし、業界をリードするモデルの FLOP 利用率を示しています。