jax

JAX - Autograd と XLA（Accelerated Linear Algebra）

JAX は、NumPy に似た API と、合成可能な関数変換の強力なシステムを組み合わせたフレームワークです。関数変換には、Grad（微分）、Jit（コンパイル）、Vmap（ベクトル化）、Pmap（並列化）があります。

使用する場面

• GPU／TPU アクセラレーションが必要な、高性能な科学シミュレーション。
• PyTorch／TF の抽象化では制約が大きすぎる、カスタムの機械学習研究。
• 最適化のための高階微分（ヘッセ行列、ヤコビ行列）の計算。
• 物理法則を組み込んだ機械学習や、微分可能なシミュレーション。
• 関数の自動ベクトル化（手動でバッチ処理する必要がなくなります）。
• コードを変更せずに、CPU、GPU、TPU で同じコードを実行する。

リファレンスドキュメント

公式ドキュメント: https://jax.readthedocs.io/
GitHub: https://github.com/google/jax
検索パターン: jax.numpy、jax.jit、jax.grad、jax.vmap、jax.random

基本原則

純粋関数（不変性）

JAX は関数型プログラミングに基づいています。すべての関数は純粋でなければなりません。つまり、（グローバル変数の変更などの）副作用がなく、同じ入力に対して常に同じ出力を返す必要があります。JAX 配列は不変です。

XLA（ジャストインタイムコンパイル）

JAX は XLA を使用して、Python／NumPy コードを特定のハードウェア向けの効率的なマシンコードにコンパイルし、最適化します。

PRNG の手動管理

NumPy とは異なり、JAX では並列環境で再現性を確保するため、乱数状態（キー）を明示的に管理する必要があります。

クイックリファレンス

インストール

# CPU
pip install jax jaxlib

# GPU (Check documentation for specific CUDA versions)
pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

標準的なインポート

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap, pmap, random

基本パターン - 微分と JIT

import jax.numpy as jnp
from jax import grad, jit

# 1. Define a pure function
def f(x):
    return jnp.sin(x) + x**2

# 2. Transform: Create a gradient function
df_dx = grad(f)

# 3. Transform: Compile for speed
f_fast = jit(f)

# 4. Use
val = f_fast(2.0)
slope = df_dx(2.0)

重要なルール

✅ 推奨事項

• jax.numpy (jnp) を使用する - NumPy と同様の機能を備えていますが、JAX の変換をサポートします。
• 純粋関数を書く - 関数が入力のみに依存し、外部状態を変更しないようにします。
• PRNG キーを手動で管理する - 乱数操作ごとに key, subkey = random.split(key) を使用します。
• バッチ処理には vmap を使用する - 単一のサンプル向けにコードを書き、バッチ次元の処理は JAX に任せます。
• JIT で static_argnums を設定する - JIT コンパイル対象の関数が配列以外の引数（ループで使用する文字列や整数など）を取る場合、その引数を static として指定します。
• .at によるインプレース更新 - 配列は不変なので、x = x.at[idx].set(val) を使用します。

❌ 禁止事項

• インプレース更新を使用する - x[idx] = val はエラーになります。
• 標準の numpy (np) を使用する - 標準の NumPy 配列は JAX の変換をサポートしません。
• 副作用を使用する - JIT コンパイルされた関数内で print() を使用したり、グローバル変数を変更したりしないでください。
• block_until_ready() を忘れる - JAX は非同期実行を使用します。正確な実行時間を測定するには、result.block_until_ready() を使用します。

アンチパターン（絶対に避ける）

import jax.numpy as jnp
from jax import jit, random

# ❌ BAD: Modifying a global variable inside a function
counter = 0
@jit
def bad_func(x):
    global counter
    counter += 1 # ❌ Side effect! Will only run once during compilation
    return x * 2

# ❌ BAD: Standard NumPy random (not reproducible/parallel-safe)
# val = np.random.randn(5) 

# ✅ GOOD: JAX PRNG
key = random.key(42)
val = random.normal(key, (5,))

# ❌ BAD: In-place assignment
# x[0] = 1.0 

# ✅ GOOD: Functional update
x = jnp.zeros(5)
x = x.at[0].set(1.0)

関数変換

Grad（微分）

def loss(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y)**2)

# Gradient of loss with respect to the 1st argument (params)
grads = grad(loss)(params, x, y)

# Higher-order: Hessian
hessian = jax.hessian(loss)(params, x, y)

Jit（ジャストインタイムコンパイル）

@jit
def complex_math(x):
    # This whole block is compiled into one XLA kernel
    y = jnp.exp(x)
    return jnp.sin(y) / jnp.sqrt(x)

# First call: Compiles (slow)
# Subsequent calls: Super fast

Vmap（自動ベクトル化）

def model(params, x):
    return jnp.dot(params, x)

# model works on 1D x. How to apply to a 2D batch of X?
# in_axes=(None, 0): don't map params, map the 0th axis of x
batch_model = vmap(model, in_axes=(None, 0))

batch_preds = batch_model(params, X_batch)

乱数（jax.random）

状態管理

key = random.key(0)

# Never reuse the same key!
key, subkey = random.split(key)
data = random.normal(subkey, (10,))

key, subkey = random.split(key)
noise = random.uniform(subkey, (10,))

PyTree の操作

複雑なデータ構造（辞書、リスト、タプル）の処理

JAX の変換は「PyTree」、つまり配列を入れ子にしたコンテナに対して機能します。

params = {'weights': jnp.ones((5,)), 'bias': 0.0}

def predict(p, x):
    return jnp.dot(x, p['weights']) + p['bias']

# grad and jit handle the dictionary automatically
grads = grad(predict)(params, x)

実践的なワークフロー

1. 微分可能な物理シミュレーション：単純な ODE を解く

def system_dynamics(state, t):
    # Simple harmonic oscillator
    x, v = state
    dxdt = v
    dvdt = -0.5 * x
    return jnp.array([dxdt, dvdt])

def loss_fn(initial_state, target_x):
    # Simulate for 10 steps using simple Euler
    state = initial_state
    dt = 0.1
    for i in range(10):
        state += system_dynamics(state, i*dt) * dt
    return (state[0] - target_x)**2

# We can take the gradient of the simulation with respect to initial state!
optimize_initial_state = grad(loss_fn)

2. vmap によるパラメータスイープ

def simulation(param):
    # Some complex computation
    return jnp.sum(jnp.linspace(0, param, 100))

# Parallelize simulation over a range of parameters
params = jnp.linspace(1, 10, 100)
results = vmap(simulation)(params)

3. pmap による分散トレーニング

# pmap replicates the function across multiple GPUs
# (assuming 8 GPUs are available)
# x = jnp.zeros((8, 1024))
# results = pmap(jnp.sin)(x)

パフォーマンスの最適化

JIT における静的引数

関数が入力値に基づくループを使用する場合、その値を static にする必要があります。

from functools import partial

@partial(jit, static_argnums=(1,))
def power_loop(x, n):
    for i in range(n):
        x = x * x
    return x

Python の制御フローを避ける

XLA の最適化を向上させるため、JAX の制御フロー（cond、while_loop、fori_loop）を使用してください。

from jax.lax import cond

def safe_divide(x, y):
    return cond(y == 0, lambda _: 0.0, lambda _: x / y, operand=None)

よくある問題と解決策

「Tracer」エラー

JIT 内では、JAX が認識するのは実際の数値ではなく、「Tracer」です。

# ❌ Problem:
# @jit
# def func(x):
#     if x > 0: return x # Error! JAX doesn't know x's value during compile

# ✅ Solution:
# Use jax.lax.cond or mark x as static_argnum

NaN 勾配

関数に（sqrt(0) のような）特異点がある場合、勾配は NaN になります。

# ✅ Solution: Add a small epsilon
def safe_sqrt(x):
    return jnp.sqrt(x + 1e-8)

GPU のメモリリーク

JAX はデフォルトで GPU メモリの 90% を事前に割り当てます。

# ✅ Solution: Set environment variable
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'

ベストプラクティス

1. 常に純粋関数を使用する - 副作用をなくし、決定論的な出力にします
2. PRNG キーを明示的に管理する - 乱数を使う操作ごとにキーを分割します
3. 頻繁に実行するループには JIT を使用する - 繰り返し呼び出される関数をコンパイルします
4. バッチ処理には vmap を活用する - 単一サンプルのコードを記述し、バッチ処理は JAX に任せます
5. 静的引数を指定する - JIT で配列以外のパラメーターには static_argnums を使用します
6. 関数型の更新を使用する - 配列を変更するときは常に .at メソッドを使用します
7. 最適化の前にプロファイリングする - jax.profiler を使用してボトルネックを特定します
8. デバイスへの配置を管理する - jax.device_put() を使用して配列の配置先を制御します
9. まず CPU でテストする - CPU でデバッグしてから、GPU／TPU に拡張します
10. コンパイルコストを理解する - JIT の初回呼び出しは遅く、2 回目以降は速くなります

JAX は、微分可能な科学を探求するための究極の遊び場です。数学を関数型変換の連続として扱うことで、これまで Python では不可能だった速度と複雑さを実現します。