JAXとGoogleのスタック

モデルと提供内容 #

どんな組織か

このハブは、新しい仕事でJAXとPyTorchのどちらを使うか選ぼうとしている人、そしてTPUの割り当てを渡されたすべての人のためのページです。JAXは2018年にGoogle Research内で、autogradとXLAの組み合わせとして始まりました。微分、ベクトル化、コンパイルができるNumPy関数です。Google DeepMindは研究の大半でこれを採用し、Googleハブで説明しているGeminiとGemmaのモデルはJAXで学習されています。Googleの以前のフレームワークであるTensorFlowは新規の研究についてはメンテナンスモードにあり、ユーザーの多くはJAX、PyTorch、Keras 3に移行しました。

この分野でのJAXのシェアはPyTorchと比べると小さいものの、集中しています。TPUユーザー、強化学習、科学計算、そして純粋関数を好む人たちです。周辺のエコシステムは「全部入り」ではなくモジュール式で、それが魅力でもあり負担でもあります。

ツールとAPI

pip install jaxでCPUビルドが入ります。GPUビルドはLinux x86_64およびaarch64向けのjax[cuda13]（またはCUDA 12版）、Cloud TPU VM向けのjax[tpu]です。インストールガイドのサポート表は正直で読む価値があります。Linux CPUとNVIDIA GPUは完全サポート、macOSはCPUのみ（Apple GPUの経路は実験的と記載され、同じページにMac GPUは非サポートとあります）、WindowsはCPUのみでWSL2経由のGPUは実験的扱い、AMD ROCmはLinux x86_64で動作、Intel GPUは実験的です。JAX 0.12.0（2026年10月19日）はNumPy 2.2以降と、Haswell世代以降のx86-64 CPUを必要とします。Python 3.11のサポートは0.11.0で終了したため、3.12以降が必要です。

JAX自体にはnn.Moduleがありません。それを提供するのがFlaxで、2024年以降に推奨されているAPIはNNXです。NNXはPyTorchと同じようにPythonオブジェクト上に状態を保持します。旧来のLinen APIは純粋関数型で、今もメンテナンスされています。Optaxがオプティマイザを、Orbaxがチェックポイントを提供し、ChexとGrainがテストとデータ読み込みを担当します。model.fitが欲しければ、これらすべての上にKeras 3が載ります。スタック全体がApache-2.0です。

中心となる考え方を数行で示します。

# snippet
import jax
import jax.numpy as jnp

def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)

grad_loss = jax.jit(jax.grad(loss))          # compile the gradient function
batched = jax.vmap(loss, in_axes=(None, 0, 0))  # per-example losses, no loop
w = jnp.zeros(4)
x = jax.random.normal(jax.random.key(0), (256, 4))
y = x @ jnp.array([1.0, -2.0, 0.5, 3.0])
for _ in range(200):
    w = w - 0.1 * grad_loss(w, x, y)
print(w)   # approaches [1, -2, 0.5, 3]

JAXはField Guideのパイプライン環境にはインストールされていないため、上のブロックは構文チェックのみ行ったスニペットです。つまずきやすい点は設計の裏返しです。jitする関数はすべて純粋でなければならず、インプレース更新はx.at[i].set(v)になり、乱数には明示的なキーが必要で、コンパイルされた関数の初回呼び出しには数秒から数分かかることがあります。XLA内部から出るエラーメッセージはPyTorchより不親切です。

TPUの料金はチップ時間単位で、世代（v5e、v5p、Trillium v6e、Ironwood v7）とリージョンによって異なります。本稿執筆時点ではGoogleの料金表を読めなかったため、Google CloudのTPU料金ページを参照してください。Hugging Face Inference EndpointsもTPU v5eのスライスを時間貸ししています。

移行メモ

PyTorchから: 大きな転換点は、パラメータが「変更する属性」ではなく「渡すデータ」になることです。PyTorchのモジュールはFlax NNXのモジュール（馴染みやすい）か、純粋関数とパラメータのpytree（より遠い）になります。optimizer.step()はtx.update(grads, state)に続けてoptax.apply_updates(params, updates)を呼ぶ形になります。データローダーは提供されないため、多くの人はtorch.utils.dataを使い続けるか、Grainを使います。当サイトのフレームワークのリファレンスに対照表があります。

モデルの重みは、どちらも単なるテンソルなので、safetensorsかNumPy配列を介して両者の間を行き来できます。Hugging Face Transformers v5はFlaxのモデルクラスを廃止したため、HubのモデルのJAX移植はFlaxAutoModelではなく別プロジェクト（MaxText、EasyDeL、またはモデル自身のリポジトリ）から得ることになります。JAXのチェックポイントをPyTorchに変換するには、パラメータ名を対応づけて線形層を転置する必要があり、退屈ですが機械的な作業です。

GPUからTPUへ: jax.shardingとshard_mapがデバイス固有の集団通信に取って代わり、同じコードが1チップでもPodスライスでも動きます。PyTorchユーザーはPyTorch/XLAを通じてTPUを使えますが、動作はするものの機能面ではネイティブのJAXに遅れています。

注意点

APIは変わり続けます。0.11と0.12のリリースノートでは、jax.custom_remat、jax.experimental.shard_alike、cloud_tpu_initモジュールが削除され、meshgridとbroadcast_arraysの戻り値の型がタプルに変わり、jnp.emptyが本当に未初期化のメモリを返すようになりました。jax、jaxlib、flaxはセットでバージョンを固定してください。

TPUはGoogle Cloudの製品であり、他では使えません。JAXで書いたコードは移植可能ですが、OrbaxがTPUシャーディングのレイアウトで書き出したチェックポイントをGPUで再読み込みするには注意が必要です。ColabとTPU Research Cloudを通じた無料のTPUアクセスは存在しますが、容量に制限があり、利用規約が付きます。

Flaxには2つのAPIがあり、ドキュメントは両方をカバーしているため、チュートリアルをコピーする前にどちらを使っているか確認してください。JAXバックエンド上のKeras 3は、Kerasユーザーにとってきれいな入口ですが、バグの追跡が難しくなる層を1つ追加します。ライブラリにテレメトリはありません。ColabとCloud TPU VMはGoogleのサービスで、それぞれ独自の規約があります。