◀ 7.

JAX とは？Google 製の関数型数値計算フレームワーク

▶

編集

本稿は JAX (ジャックス) に関する記事です。

この記事の要点

JAX は "NumPy 互換 API + 自動微分 + GPU/TPU + JIT" を統合した数値計算フレームワーク
Google が開発・Google DeepMind の研究の主軸。Gemini など先端モデルの学習でも採用
関数型のスタイル（純関数 + 変換）が独特。慣れが必要だが性能・並列化が強力
jit / grad / vmap / pmap という関数変換で書く
JAX 単体は低レベル。実用ではニューラルネット用上位ライブラリ Flax / Equinox を併用
TPU を活かせる事実上唯一の主要フレームワーク（PyTorch も対応進行中）

JAX とは？

JAX は、Google が開発する数値計算ライブラリです。「NumPy のような API」「自動微分」「GPU / TPU での実行」「XLA による JIT コンパイル」「並列化変換」を、関数を変換する形で統一的に提供します。

PyTorch / TensorFlow が「フレームワーク」だとすれば、JAX は「NumPy の強化版」に近い思想で設計されており、ニューラルネット以外の数値最適化・物理シミュレーション・確率モデル等にも自然に使えます。深層学習向けの上位ライブラリ (Flax / Equinox / Haiku) を上に積んで使うのが一般的です。

主な特徴

特徴	説明
NumPy 互換 API	`jax.numpy as jnp` がほぼ NumPy と同じ呼び方
自動微分 (`grad`)	任意関数の偏微分・高階微分を関数変換として取得
JIT コンパイル (`jit`)	XLA で関数全体を最適化。GPU/TPU でも高速
ベクトル化 (`vmap`)	ループ不要でバッチ並列化
デバイス並列 (`pmap` / `pjit`)	複数 GPU / TPU 跨ぎの並列実行
純関数モデル	副作用なしの関数を変換していくスタイル
TPU 親和性	Google TPU で性能が最も出やすい
関数型エコシステム	Flax / Optax / Orbax / Chex 等のモジュール群

4 つの主要変換

変換	役割	直感
`jit(fn)`	関数を XLA でコンパイル	「同じ計算を爆速にする」
`grad(fn)`	勾配を返す関数を生成	「微分してくれる」
`vmap(fn)`	関数を自動でバッチ化	「for ループを並列化」
`pmap(fn)` / `pjit`	複数デバイス間で並列実行	「複数 GPU/TPU に分散」

最小サンプル: NumPy 風＋自動微分＋JIT

import jax
import jax.numpy as jnp

# NumPy 風に書ける
def f(x):
return jnp.sum(jnp.sin(x) ** 2)

# 勾配を取る関数を生成
grad_f = jax.grad(f)

# JIT コンパイル
fast_grad_f = jax.jit(grad_f)

x = jnp.array([0.0, 1.0, 2.0])
print(fast_grad_f(x))

関数型スタイルの注意点

PyTorch / NumPy 出身者が戸惑うポイント

配列はイミュータブル: x[0] = 1 はできない。x = x.at[0].set(1) を使う
純関数でないと jit / vmap が破綻する (外部状態・Python の制御フロー注意)
Python の for 内で副作用を出すと、JIT 時にトレースが効かない
ランダム数は明示的な PRNG キー (jax.random.PRNGKey) を渡す
シェイプ依存の分岐は jax.lax.cond / jax.lax.fori_loop を使う
NaN や形状不一致はコンパイル後に出るので、jax.config.update("jax_disable_jit", True) でデバッグ

深層学習用上位ライブラリ

ライブラリ	立ち位置
Flax (Linen / NNX)	Google 公式。Gemini 等の社内利用実績。現在は NNX が新世代
Equinox	PyTorch ライクな素直さ。クラスで層を書く
Haiku	DeepMind 製。Flax の前身的存在 (現在は保守モード)
Optax	最適化アルゴリズム (Adam / AdamW / Lion 等)
Orbax	チェックポイント保存・配信
Chex	テスト・アサーションのユーティリティ
NumPyro	確率プログラミング (Pyro の JAX 版)
Diffrax	微分方程式ソルバ
jaxtyping	テンソルの形状・型を注釈

PyTorch / TensorFlow との立ち位置

観点	JAX	PyTorch	TensorFlow
パラダイム	関数型	命令型 (eager)	命令型 (eager) + Keras
TPU 親和性	最強	対応進行中	従来からの定番
研究での採用	Google 系・確率プログラミング	圧倒的多数	減少傾向
本番デプロイ	少ない (JAX2TF 経由)	TorchServe / ONNX	TF Serving / TFLite / TF.js
学習コスト	高め (関数型 + 純関数)	低い	中
科学計算・物理	得意	○	○

使うのが向くシーン

TPU を活用したい（Google Cloud TPU、Colab TPU）
大規模分散学習を関数変換で表現したい
確率プログラミング・科学計算・物理シミュレーション
研究プロトで論文の実装を NumPy 風に純粋に書きたい
Google DeepMind / Google Brain 系の論文を再現したい

インストール

# CPU 版
pip install jax

# NVIDIA GPU 版 (CUDA 12)
pip install -U "jax[cuda12]"

# TPU (Google Cloud / Colab TPU)
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# よく一緒に入れるもの
pip install flax optax

運用上のヒント

Tips

最初は jit 無しで書いて挙動を確認 → 安定したら jit を被せる
vmap を使うクセ: バッチ次元のループを書きたくなったらまず vmap
勾配が大きい / NaN が出る場合は jax.debug.nan_check や jax.debug.print
大規模モデルは FSDP 相当の pjit / shard_map でテンソル分割
チェックポイントは Orbax。学習中断・再開や非同期保存
PyTorch コードを移植する場合、まず NumPy 部分を jnp に置換 → jit 化のステップを踏む

注意点

よくある落とし穴

初回 JIT コンパイルが遅い: 起動時のキャッシュコストを許容できる用途を選ぶ
動的な形状でループするとコンパイルし直しになる。形状を固定化する設計
イミュータブル配列を理解しないと初学者が詰まる
本番デプロイのエコシステムが PyTorch / TF より薄い
GPU でのデバイス・メモリ管理 (XLA_PYTHON_CLIENT_PREALLOCATE 等) が独特
Windows ネイティブの GPU サポートは限定的。Linux / WSL2 を推奨
JAX とそのエコシステム (Flax / Optax) はバージョン整合が崩れやすい。lock ファイルで固定

ページ一覧

その他