jax-ml/jax-ai-stack

これは何？

主な機能

• 01JAXと、Flax、Optax、Orbax、Chex、Grain、ml_dtypesを、組み合わせてテスト済みのバージョンでインストールできます。
• 02jit、vmap、gradなどのJAX変換機能を利用できます。
• 03pip経由でtensorflow-datasetsをオプションとして追加できます。
• 04CUDA GPUやTPU向けのハードウェア固有のJAXインストールに対応しています。

ユースケース

• FlaxとOptaxを使ったニューラルネットワークの学習環境をJAXで構築する。
• JAXプロジェクトでOrbaxを使い、モデルのチェックポイントを保存・復元する。
• 動作確認済みの依存関係セットを使い、NVIDIA GPUやGoogle TPUでJAXの処理を実行する。

クイックスタート

1. 1

   PyPI から JAX AI スタックをインストールする

   pip install jax-ai-stack
2. 2

   必要に応じて、tfds extra を追加して TensorFlow Datasets のサポートを有効にする

   pip install jax-ai-stack[tfds]
3. 3

   必要に応じて、同じコマンドに JAX CUDA extra を追加して GPU 対応でインストールする

   pip install jax-ai-stack "jax[cuda]"

使う前に知っておきたいこと

• !他のライブラリをまとめたバンドルのため、固定されたパッケージが提供する機能以外に、単独で利用できる機能はありません。
• !README によると、ドキュメントとチュートリアルはまだ作成中です。

収録されている特集

トピック

よくある質問

jax-ai-stack とは？

AIや数値計算向けに、テスト済みのJAXライブラリ（Flax、Optax、Orbax、Chex、Grain）を互換性のあるバージョンでまとめてインストールできるパッケージ。 JAX AI Stackは、連携して動作するJAXベースのライブラリ群をまとめてインストールするためのエントリーポイントです。`pip install jax-ai-stack`を実行すると、JAXに加えて、ニューラルネットワーク構築用のFlax、勾配処理と最適化用のOptax、チェックポイント管理用のOrbax、信頼性の高いJAXコードを書くためのChex、データ読み込み用のGrain、NumPyのdtype拡張を提供するml_dtypesなどのコンポーネントがインストールされます。各バージョンは、リポジトリの統合テストで検証済みの組み合わせに固定されています。

jax-ai-stack は無料で使えますか？ライセンスは？

jax-ai-stack は Apache-2.0 ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。

jax-ai-stack は今もメンテナンスされていますか？

活発にメンテナンス中。jax-ai-stack の最終コミットは 2026年10月13日 です。

jax-ai-stack は何に使えますか？

FlaxとOptaxを使ったニューラルネットワークの学習環境をJAXで構築する。 JAXプロジェクトでOrbaxを使い、モデルのチェックポイントを保存・復元する。 動作確認済みの依存関係セットを使い、NVIDIA GPUやGoogle TPUでJAXの処理を実行する。

jax-ai-stack のインストール方法は？

PyPI から JAX AI スタックをインストールする: pip install jax-ai-stack 必要に応じて、tfds extra を追加して TensorFlow Datasets のサポートを有効にする: pip install jax-ai-stack[tfds] 必要に応じて、同じコマンドに JAX CUDA extra を追加して GPU 対応でインストールする: pip install jax-ai-stack "jax[cuda]"

jax-ai-stack はどこで入手できますか？

ソースコードとインストール手順は GitHub にあります：https://github.com/jax-ml/jax-ai-stack