JAX

どんなツールか

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

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

横にスクロール

特徴・仕組み

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

• grad: 関数を微分した関数を返す（自動微分。高階微分も合成で得られる）。
• jit: XLA コンパイラで関数をコンパイルし、カーネル融合などで高速化する。
• vmap: 処理を自動でベクトル化（バッチ化）し、ループを書かずに一括処理する。
• pmap / シャーディング: 複数の GPU/TPU へ計算を分散する。

NumPy 互換の API を持ちつつ、これらの変換を自由に組み合わせて性能を引き出すのが基本です。内部では計算を一度トレースして中間表現（jaxpr）にし、XLA が対象ハードウェア向けに最適化します。

得意・不得意

JAX は 関数型のスタイルを前提とします。状態を持たない純粋関数・不変な配列として処理を書くことで、変換が安全に適用できる設計です。乱数も明示的な鍵（key）で管理し、再現性を担保します。この性質から、GPU / TPU を活かした高性能計算や大規模研究で台頭してきました。

一方で、書き方の作法（純粋関数・不変値・明示的乱数）に慣れが要り、jit 対象では動的な配列形状や Python の副作用が扱いにくい制約があります。学習・配信まわりの標準は薄めで、ニューラルネット記述は Flax / Haiku、最適化は Optax といった上位ライブラリを組み合わせて使うのが一般的です。

PyTorch との違い

使いどころ・注意点

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

総じて JAX は、NumPy 的な書き味と合成可能な変換（grad/jit/vmap）で高性能計算を実現する、研究・数値計算志向の強力なライブラリです。