JAXとは　関数を変換して速くするAIの道具、TPUの9216チップ構成も視野

JAXは、AIや科学技術の計算を書くためのPythonのライブラリです。公式の説明は「アクセラレータ向けの配列計算とプログラム変換のためのライブラリ」です。グーグルの研究者が始め、今はjax-mlという公開の場で開発が続いています。

関数を「変換」する

JAXの特徴は、ふつうのPythonの関数を別の働きを持つ関数に作り替える点にあります。数値の配列の扱いは、NumPyという定番の道具とほぼ同じ書き方です。

JAXの主な変換（JAXの公式文書から）

変換	働き
jax.jit	関数をまとめてコンパイルし、半導体向けの速い命令にする
jax.grad	関数の勾配（微分）を自動で計算する
jax.vmap	1件用に書いた関数を、多数の件をまとめて処理する形にする
jax.shard_map	計算を複数の半導体に分け、それぞれで動かす

jitで変換すると、関数はXLAというコンパイラに渡り、CPU、GPU、TPUそれぞれ向けの命令になります。同じコードのまま、動かす半導体を替えられます。PyTorchが書いた順にその場で計算を進めるのに対し、JAXは関数全体をまとめて最適化する考え方です。

道具の一式と版の進み方

JAXの上には、ニューラルネットワークを組むFlax、学習の最適化を担うOptax、学習途中の状態を保存するOrbax、大規模言語モデル用のMaxTextなどが載ります。いずれも公式文書が周りの道具として挙げています。

変更履歴によると、0.11.2版が2026年9月17日、0.12.0版が10月19日、0.12.0.1版が10月23日に出ました。0.11.0版（7月16日）ではPython 3.11への対応をやめています。0.12.0版では、TPUを動かす部品（libtpu）の最低の版を0.0.50に上げました。AMDの開発基盤ROCmの10.1版も、JAX 0.11.1を対応の一覧に入れています。

TPUとの組み合わせ

JAXが最も多く使われる場面の一つが、グーグルのTPUです。グーグルが2025年4月9日に発表した第7世代のTPU「Ironwood」は、1チップに192GBのHBMを積み、メモリの速さは毎秒7.37テラバイトです。最大9,216チップを1つの塊（ポッド）につなぎ、合計の計算性能は42.5エクサFLOPSになります。

グーグルは、何万個ものTPUをまとめて使うための実行の仕組み「Pathways」も外部に開くと説明しました。Anthropicは2025年10月23日、グーグルのTPUを最大100万個使う契約を発表しています。

何千個もの半導体に計算を分けるには、データやモデルをどう切り分けるかを書く必要があります。JAXでは、配列をどの半導体に置くかを指定し、shard_mapなどで分割の仕方を決めます。

関連する用語

• XLA：JAXやTensorFlowが使うコンパイラ。計算をまとめて半導体向けの命令にする。
• TPU：グーグルがAIの計算向けに設計した半導体。
• Flax：JAXでニューラルネットワークを組むためのライブラリ。
• PyTorch：メタが育てた、もう一つの代表的な枠組み。

グーグルは2026年4月、第8世代のTPUを学習用の「TPU 8t」と推論用の「TPU 8i」に分けると発表しています。