Flax

Identity

• FlaxはJAX向けの柔軟性を重視したニューラルネットワークライブラリで、Google DeepMindのチームが維持するオープンソースプロジェクトである（公式Googleプロダクトではない）。検証待ち GitHub - google/flax

Capability

• FlaxはJAX向けの柔軟性を重視したニューラルネットワークライブラリで、Google DeepMindのチームが維持するオープンソースプロジェクトである（公式Googleプロダクトではない）。検証待ち GitHub - google/flax
• Flaxはflax.nnxというニューラルネットワークAPIを提供し、Linear・Conv・BatchNorm・LayerNorm・Attention（MultiHeadAttention）・LSTMCell・GRUCell・Dropout等のコンポーネントと、レプリケート学習・シリアライズ・チェックポイント等のユーティリティを含む。検証待ち Neural networks with JAX - Flax
• Google AI for Developersの公式ページに『Flax was used to create Gemma』と明記されており、RecurrentGemma（Griffinモデル）のライブラリもJAX・Flax・Orbax・SentencePieceで実装されている。GriffinモデルはFlaxのflax.linen.Moduleを継承する構造を持つ。検証待ち Fine-tuning RecurrentGemma using JAX and Flax（Google AI for Developers）

Credibility

• Flaxは、既存の学習ループを変更するのではなくフォークして新しい学習手法を試すという設計思想を持つ点で、他のNNフレームワークとは異なるアプローチを取っている。検証待ち flax/README.md at main · google/flax

Use Case

公開Evidence未整備

Constraints / Current Status

• Flaxはflax.nnxというニューラルネットワークAPIを提供し、Linear・Conv・BatchNorm・LayerNorm・Attention（MultiHeadAttention）・LSTMCell・GRUCell・Dropout等のコンポーネントと、レプリケート学習・シリアライズ・チェックポイント等のユーティリティを含む。検証待ち Neural networks with JAX - Flax
• Google DeepMindが管理するHaiku（別のJAX向けニューラルネットワークライブラリ）のGitHubリポジトリのREADMEに、2023年7月時点でGoogle DeepMindが新規プロジェクトへFlaxの採用を推奨するとの記述がある。理由として開発チーム規模、機能の豊富さ、Alphabet外での採用実績の多さが挙げられ、Haikuはバグ修正・JAX互換性維持中心のメンテナンスモードへ移行するとされている。検証待ち google-deepmind/dm-haiku README（Flaxへの移行推奨に関する記述）

Key References

• Google DeepMindはなぜFlaxを推奨しているのですか？
• Flaxとはどのようなライブラリですか？
• Flaxは他のJAX向けライブラリと何が違いますか？

Part of

• Google DeepMind— FlaxはGoogle DeepMindのチームが維持するライブラリである。

Related

• JAX— FlaxはJAX上に構築されたニューラルネットワークライブラリであり、JAXと直接関連する。