Skip to main content

JAX 日本語チュートリアル

このドキュメントは、高性能な数値計算と機械学習のための Python ライブラリ JAX を、Python にあまり詳しくない方でも理解できるように、日本語で丁寧に解説するチュートリアルです。

JAX は「NumPy ライクな API で書ける」「自動微分ができる」「JIT コンパイルで高速化できる」「CPU/GPU/TPU で同じコードが動く」「ベクトル化・分散並列化を関数変換として扱える」という特徴を持つ、ちょっと風変わりで強力なライブラリです。Google の研究部門や DeepMind、その他多くの研究機関が、大規模な機械学習や科学計算の現場で JAX を使っています。

本チュートリアルは、JAX の 公式ドキュメント、JAX の GitHub リポジトリ、Google Cloud Blog の PyTorch 開発者向け JAX ガイド（日本語）、Zenn の Turing Motors による JAX/Flax 入門記事 を参考に、2026 年時点の最新情報に合わせて整理しました。

このチュートリアルの想定読者​

  • Python の基本文法（変数・関数・リスト・for 文）はなんとなく分かる
  • NumPy を少し触ったことがある、または「なんか聞いたことがある」レベル
  • 機械学習・ディープラーニングに興味がある
  • PyTorch や TensorFlow の経験は なくても OK

「Python は学校で少しやった」「コードはまだ書き慣れていない」という方でも、できるだけ専門用語に頼らず、具体例を交えて読み進められるように書いています。読み始める前に、まず 用語集 に軽く目を通すと、本編がぐっと分かりやすくなります。

全体構成​

本チュートリアルは 基礎編（第1〜12章） と 移行ケース集（第13〜15章） の 2 部構成です。

  • 基礎編：JAX という言語そのものに入門するチュートリアル。順番に読み進めると JAX で書ける状態になります。
  • 移行ケース集：すでに NumPy や PyTorch で書かれたコードを JAX に書き換えるためのレシピ集。「左に元のコード、右に JAX 版」の対比形式で、目的別にすぐ引けます。基礎編をひととおり読んだ後に辞書のように使うのもおすすめです。

目次​

📘 基礎編：JAX の入門​

  1. 用語集（最初に読むと便利）
  2. JAX とは何か
  3. インストールと環境構築
  4. jax.numpy の基本
  5. JIT コンパイル（jit）で高速化する
  6. 自動微分（grad）
  7. ベクトル化（vmap）と並列化（pmap / shard_map）
  8. PyTree という考え方
  9. 乱数の扱い方（PRNG キー）
  10. PyTorch との比較
  11. 実践チュートリアル：線形回帰と簡単なニューラルネット
  12. よくある落とし穴とデバッグ
  13. JAX エコシステム（Flax / Optax / Equinox など）

🔁 移行ケース集：既存コードを JAX に書き換える​

  1. NumPy → JAX：配列・乱数・ループ・線形代数・FFT・SciPy など、NumPy の典型パターンを JAX で書く 15 ケース
  2. PyTorch → JAX：テンソル操作・モデル定義・学習ループ・損失/最適化・CNN/RNN/Attention・分散学習など、PyTorch の典型パターンを JAX（+ Flax / Optax）で書く 20 ケース
  3. 応用と複合パターン：アンサンブル学習・メタ学習・ベイズ推論・強化学習・微分可能シミュレーション・K-Means・ニュートン法など、「JAX らしい書き方」が活きる 12 ケース

🧬 進化計算（Evolutionary Computation）​

  1. 進化計算入門と基本アルゴリズム：ランダム探索 / (1+1)-ES / GA / DE / PSO / OpenAI-ES / CMA-ES / MAP-Elites / NSGA-II
  2. ニューロエボリューションと進化的モデルマージ：MLP/方策ネットの ES 学習、CartPole の RL、evosax の使い方、Sakana AI の Evolutionary Model Merge、勾配 × 進化のハイブリッド

読み進め方​

  • はじめての方：用語集 → 第1章 → 第2章 …と 順番に 読むのがおすすめです。第10章で実際に手を動かし、第11章で「ハマりやすいポイント」を確認するところまで到達すると、JAX の世界がぐっと身近になります。その後、必要に応じて第13〜15章のケース集を辞書代わりに参照してください。
  • PyTorch ユーザの方：先に第9章「PyTorch との比較」をざっと眺め、第3〜8章で JAX 独自の流儀を抑えてから、第14章「PyTorch → JAX 移行ケース集」 に進むのが最短ルートです。
  • NumPy / SciPy ユーザの方：第1章 → 第3章 → 第4章 → 第5章 →（必要なら第8章）→ 第13章「NumPy → JAX 移行ケース集」 の順がスムーズです。
  • すでに JAX を触ったことがある方：気になる章だけ拾い読みしても大丈夫です。各章は独立して読めるように書いてあります。

目的別クイックリンク（移行ケース集の早見表）​

やりたいこと行き先
np.xxx の JAX 版を知りたい13.1〜13.14 と 13.15 小技対応表
x[i] = v を JAX でどう書くか13.2 in-place 更新
for ループを高速化したい13.5 累積計算 → scan
np.random.seed から JAX 乱数へ13.6 乱数生成、第8章
数値微分を jax.grad に置き換えたい13.10 自動微分
PyTorch の学習ループを移植したい14.5 学習ループ全体
nn.Module を JAX でどう書くか14.4 線形層（MLP）、14.11 CNN
loss.backward() の置き換え14.3 自動微分
Adam / AdamW / Scheduler の対応14.7・14.8
Dropout / BatchNorm の扱い14.10
RNN / Transformer の移植14.12 RNN、14.13 Attention
DDP / FSDP 相当の分散学習14.17 マルチ GPU / 分散学習
detach / gather / scatter などの対応14.20 小技対応表
アンサンブルを vmap で並列化15.2 アンサンブル学習
ハイパラ探索を並列で回す15.3 ハイパラ並列探索
MAML 風メタ学習15.4 メタ学習
強化学習の並列ロールアウト15.6 強化学習
微分可能シミュレーション15.7 物理シミュレーション
進化計算を JAX で書きたい第16章 進化計算入門
GA / ES / CMA-ES の実装例16.5 GA / 16.8 OpenAI-ES / 16.9 CMA-ES
ニューラルネットを進化させたい第17章 ニューロエボリューション
強化学習を進化計算で（CartPole）17.4 RL タスクへの応用
進化的モデルマージ（Sakana AI）17.6 進化的モデルマージ

参考にした主な資料​

💡 各章の末尾には「次の章へのリンク」が付いています。気軽にコードをコピーして、手元で試しながら読み進めてみてください。動かして初めて分かることがたくさんあります。