1 歴史と背景
1.1 開発の経緯
JAXは、Googleの研究チームによって、機械学習と数値計算のための統一フレームワークとして開発された。従来のTensorFlowが持つ静的グラフモデルの制約を克服し、動的な計算グラフとNumPyとの高い互換性を実現することを目的とした。開発の初期段階では、自動微分ライブラリ「Autograd」と線形代数コンパイラ「XLA（Accelerated Linear Algebra）」を組み合わせるアイデアが中心となり、研究者がより直感的に高速な実験を行える環境を提供することを目指した。
1.2 公開とバージョン履歴
JAXは2018年12月にオープンソースとして初めて公開された。初期バージョンでは基本的なNumPy互換関数と自動微分機能が提供され、その後JITコンパイルや並列化機能が段階的に追加された。主要なバージョンとしては、2020年のv0.2シリーズでpmapやvmapの安定化、2022年のv0.3シリーズでWindowsサポートの拡充やパフォーマンス最適化が行われた。2024年現在も継続的に開発が進められ、最新版では複数のバックエンド（CPU/GPU/TPU）に対する高度な最適化が実現されている。
2 コア機能
2.1 自動微分（grad）
JAXのgrad関数は、任意のPython関数を入力として受け取り、その関数の勾配を計算する関数を返す。自動微分はリバースモード（逆伝播）を基本とし、高階微分や複数引数に対する微分もサポートする。NumPyライクなコードをそのまま微分可能に変換できるため、機械学習の損失関数や科学計算の感度解析に広く利用される。
2.2 JITコンパイル（jit）
jit関数は、Python関数をXLA（Accelerated Linear Algebra）を用いてJust-In-Timeコンパイルし、実行速度を大幅に向上させる。関数内の演算を静的に解析し、ハードウェアに最適化された低レベルコードに変換することで、特に大規模な配列演算やループ処理で顕著な高速化を実現する。jitはトレースベースのコンパイルを採用し、形状やデータ型が固定された関数に対して効率的なコード生成を行う。
2.3 ベクトル化（vmap）
vmapは、関数の自動ベクトル化（ベクトル化写像）を提供する変換である。バッチ次元を自動的に処理し、ループを明示的に書かずに複数の入力サンプルに対して一括処理を実行できる。例えば、単一データ点に対する関数をvmapでラップすることで、バッチ全体の処理を効率的にベクトル化された演算に変換する。これにより、コードの簡潔さとパフォーマンスの両立が可能となる。
2.4 並列化（pmap）
pmapは、複数のデバイス（GPUやTPUコア）にまたがるデータ並列実行を実現する関数変換である。ユーザーは単一デバイスで動作する関数を記述するだけで、pmapが自動的にデータ分割と各デバイスへの割り当て、結果の収集を行う。大規模な分散学習や科学研究における大規模シミュレーションで活用され、線形なスケーリング性能を達成することが可能である。
2.5 乱数生成（PRNG）
JAXは、従来の状態を持つ乱数生成器ではなく、分割可能な疑似乱数生成器（PRNG）を採用している。jax.random.PRNGKeyでシードを生成し、split関数で独立したサブキーを作成することで、再現性と並列化時の安全性を保証する。この設計により、複数のスレッドやデバイスで異なる乱数列を生成する際の競合を回避し、関数型プログラミングの原則に適合する。
3 エコシステムと主要ライブラリ
3.1 Flax
Flaxは、JAXをベースとした高レベルなニューラルネットワークライブラリである。Google Researchが中心となって開発され、モデル定義のためのlinenモジュールを提供する。動的なモジュールシステム、柔軟な訓練ループの設計、事前学習モデルのコレクションなどを備え、研究用途に特化した設計が特徴。Optax（最適化）、Orbax（チェックポイント）、TensorStore（データ管理）などJAXエコシステムと密接に連携する。
3.2 Haiku
Haikuは、DeepMindが開発したJAX向けのニューラルネットワークライブラリである。Sonar（Sonnet）の設計思想を継承し、関数型のAPIを提供する。パラメータ管理は明示的で、Pythonの関数内で定義された層のパラメータを自動的に収集する仕組みを持つ。特に強化学習や大規模なモデル研究で使用され、DeepMindの多くのプロジェクトで採用された実績がある。
3.3 Optax
Optaxは、JAX用の勾配ベースの最適化ライブラリである。SGD、Adam、AdamW、Lambなどの一般的な最適化手法に加え、勾配クリッピング、スケジューリング、損失スケーリングなどのユーティリティを提供する。関数型の設計により、複数の最適化手法の組み合わせやカスタム勾配変換を容易に実装できる。JAXエコシステムの事実上の標準最適化ライブラリとして広く利用されている。
3.4 DeepMindとの連携
DeepMindはJAXの主要なユーザーかつコントリビューターであり、複数の研究プロジェクトやライブラリにおいてJAXを基盤として採用している。例えば、AlphaFold2の一部の処理、強化学習フレームワークのAcme、生成モデルライブラリのDopamineなどがJAX上で実装されている。また、HaikuやOptax、rlax（強化学習ユーティリティ）などのライブラリを公開し、研究コミュニティへの貢献を行っている。
4 使用例と応用分野
4.1 機械学習研究
4.1.1 強化学習
JAXは強化学習研究において、その高速なシミュレーションとエピソード単位の並列実行が評価されている。vmapとpmapを活用して環境インスタンスを並列化し、経験収集の効率を向上させる。また、gradを用いた方策勾配法の実装や、Optaxによる安定した最適化が可能となる。DeepMindのAcmeフレームワークや、シンプルなアルゴリズム実装としてのPureJaxRLなど、多くのプロジェクトで利用されている。
4.1.2 生成モデル
拡散モデルや変分オートエンコーダ、GANなどの生成モデルの実装においてもJAXは広く採用されている。JITコンパイルによる高速なサンプリングと、vmapによるバッチ生成の効率化が研究の生産性を向上させる。FlaxやHaikuを用いたモデル定義と、Optaxによる訓練の組み合わせが標準的であり、Hugging FaceのDiffusersライブラリでもJAXバックエンドが提供されている。
4.2 科学計算
4.2.1 微分方程式の解法
JAXの自動微分とJITコンパイルは、常微分方程式や偏微分方程式の数値解法に応用される。特に、ニューラル常微分方程式（Neural ODE）の実装では、勾配計算を伴うソルバーの効率化に貢献する。また、物理情報を損失関数に組み込む物理インフォームドニューラルネットワーク（PINN）の実装でも、gradによる高階微分の計算が活用される。
4.2.2 ベイズ推論
ベイズ推論においては、MCMC（マルコフ連鎖モンテカルロ法）や変分推論の実装にJAXが利用される。vmapによる独立なサンプル系列の並列実行や、gradによる勾配ベースの変分推論（ADVI）の高速化が可能となる。また、確率プログラミングライブラリNumPyroはJAXをバックエンドとし、柔軟なモデル定義とサンプリングを提供している。
5 コミュニティと影響
5.1 学術界での採用
JAXは機械学習のトップ会議（NeurIPS, ICML, ICLRなど）の論文において、主要な実装フレームワークとして頻繁に引用されている。特に強化学習、生成モデル、科学計算の分野でその採用が顕著であり、2023年以降はPyTorchに次ぐ存在感を示している。また、オープンソースの研究プロジェクトやリポジトリ（例えば、Google ResearchのJAXプロジェクトや、GitHub上のAwesome JAXリスト）で多くのリソースが公開されている。
5.2 批判と限界
JAXの主な批判点として、学習曲線の急峻さが挙げられる。関数型プログラミングのスタイルや、状態管理の徹底（例：PRNGの明示的な分割）は、オブジェクト指向に慣れた開発者には難解に映る。また、TensorFlowやPyTorchと比較して高レベルAPIの整備が遅れており、エコシステムの成熟度では劣る部分がある。さらに、動的制御フロー（条件分岐やループ）の取扱いが制限される場合があり、複雑なモデルの実装に工夫を要する。
5.3 類似フレームワークとの比較
5.3.1 PyTorch
PyTorchは動的計算グラフとオブジェクト指向のAPIを特徴とし、最も広く使われている深層学習フレームワークである。JAXと比較すると、直感的なコード記述が可能で学習が容易な一方、純粋関数型ではなく状態管理が暗黙的である。パフォーマンス面では、JAXのJITコンパイルが大規模バッチや多次元配列で優位に立つことが多いが、PyTorchもTorchScriptやtorch.compileで追従している。
5.3.2 TensorFlow
TensorFlowは静的グラフと動的グラフ（Eager Execution）の両方をサポートする。JAXはNumPyとの互換性を重視し、より細かい変換（jit, grad, vmap）を提供する。TensorFlowは産業用デプロイや分散システムとの統合で強みを持つが、研究用途ではJAXの柔軟性が評価される。また、TensorFlowの高レベルAPI（Keras）はJAXのエコシステムより充実している。
5.3.3 Julia
Juliaは数値計算と科学技術計算に特化した言語であり、JAXと同様に自動微分（Zygote.jlなど）やJITコンパイル（LLVMベース）を自然に扱う。Juliaのアドバンテージは、多重ディスパッチによる表現力と、純粋な言語機能としての自動微分である。一方、JAXはPythonエコシステムとの親和性が高く、機械学習コミュニティの資産を活用できる点で有利である。
6 関連項目
- 自動微分
- XLA（Accelerated Linear Algebra）
- NumPy
- 機械学習フレームワーク
- Google Research
- DeepMind