JAX

どんなツールか

JAX は Google が公開した数値計算ライブラリです。オープンソースで、高性能な機械学習や科学技術計算を念頭に設計されています。

「結局なに？」を一言でいえば、NumPy に近い書き味 のまま、自動微分とコンパイルによる高速化を足し、GPU / TPU で動かせるようにしたライブラリです。

特徴・仕組み

JAX の中心は、関数に対して適用する一連の 変換 です。

• grad：関数を微分した関数を返す（自動微分）。
• jit：XLA コンパイラで関数をコンパイルして高速化する。
• vmap：処理を自動でベクトル化（バッチ化）する。

NumPy 互換の API を持ちつつ、これらの変換を組み合わせて性能を引き出すのが基本的な使い方です。

得意・不得意

JAX は 関数型のスタイル を前提とします。状態を持たない純粋関数として処理を書くことで、変換が安全に適用できる設計です。この性質から、GPU / TPU を活かした高性能計算や大規模な研究で台頭してきました。

一方で、書き方の作法（純粋関数・不変な値の扱いなど）に慣れが要り、学習・配信まわりの周辺ライブラリのエコシステムは PyTorch より小さいのが実情です。

いつ使うか（他との違い）

高速化やスケールを突き詰めたい研究・数値計算、TPU を活用したい場面で有力です。NumPy の延長で書けるため、数式に近い実装をそのまま高速に回したい用途にも向きます。

幅広いサンプルや学習済みモデルの蓄積を重視するなら PyTorch、性能と関数型の表現力を重視するなら JAX、と切り分けると選びやすくなります。