JAX
詳細情報
- 作成者
- 発表
- 2018年12月
- 管理元
- 種別
- 機械学習フレームワーク
- ライセンス
- Apache 2.0
- 公式サイト
- jax.readthedocs.io
- ソースコード
- github.com/jax-ml/jax
GoogleのMLフレームワーク。NumPy互換API、自動微分、XLAによるJITコンパイル、そしてTPUの第一級サポートを備える。
関数型のパラダイムで、可変状態を持たず、モデルは(params, x) → yという関数になる。主役の変換は四つ。jit（コンパイル）、grad（自動微分）、vmap（自動バッチ化）、pmap／shard_map（複数デバイス）。研究コードをきわめて簡潔に書ける。
Google DeepMind（GeminiやAlphaFold）、AnthropicのClaudeの学習、Google社内のML用途の大部分を支えている。業界全体ではPyTorchが優勢だが、TPUで作業するならJAXはほぼ必須である。