JAX
JAX は「ステロイド入りの NumPy」です。Autograd（自動微分）と XLA（コンパイル）を組み合わせます。2025 年には Flax NNX（PyTorch スタイルの OOP）が標準になりつつあります。
使用すべきケース
- 高性能数値・科学計算: GPU/TPU 上での高速線形代数のための Autograd と XLA コンパイル。
- カスタム深層学習研究: フレームワークの肥大化を避け、Flax、Haiku、Equinox でニューラルネットワークを構築。
- 構成可能な関数変換:
jit（コンパイル）、grad（微分）、vmap（ベクトル化）をシームレスに連鎖。 - マルチデバイス分散トレーニング:
pmapとシャードマップ (jax.experimental.shard_map) を使用してクラスタ全体で計算を並列化。
クイックスタート
import jax
import jax.numpy as jnp
# JIT コンパイルされた行列積と自動微分
@jax.jit
def loss_fn(w, x, y):
pred = jnp.dot(x, w)
return jnp.mean((pred - y) ** 2)
# 値と勾配を同時に計算
grad_fn = jax.grad(loss_fn)
w = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [2.0, 1.0]])
y = jnp.array([2.0, 4.0])
grads = grad_fn(w, x, y)
print("Gradients:", grads)
コアコンセプト
#構成可能な関数変換: jit, grad & vmap
ジャストインタイムコンパイル、自動微分、自動バッチ処理の組み合わせ:
import jax
import jax.numpy as jnp
# 純粋関数の定義
def loss_fn(w, x, y):
pred = jnp.dot(x, w)
return jnp.mean((pred - y) ** 2)
# 変換 1: 重みに関する自動勾配
grad_fn = jax.grad(loss_fn)
# 変換 2: XLA による JIT コンパイルで GPU を高速化
fast_grad_fn = jax.jit(grad_fn)
# 変換 3: バッチ全体の自動ベクトル化
# batch_loss_fn はループなしで 2D 行列を自動処理
vmapped_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))
# サンプル計算
key = jax.random.PRNGKey(42)
w = jax.random.normal(key, (3,))
x = jnp.array([[1.0, 2.0, 3.0]])
y = jnp.array([5.0])
grad_val = fast_grad_fn(w, x, y)
print("Computed Gradient:", grad_val)
#明示的な PRNG キー管理
状態を持たない疑似乱数生成:
import jax
# ルート PRNG キーを初期化
key = jax.random.PRNGKey(2026)
# JAX キーは明示的に分割する必要があり、再利用は不可
key, subkey1, subkey2 = jax.random.split(key, 3)
data_normal = jax.random.normal(subkey1, shape=(4, 4))
data_uniform = jax.random.uniform(subkey2, shape=(4, 4))
print("Normal sample mean:", jnp.mean(data_normal))
#Equinox と純粋関数によるステートフルモデル
JAX の純粋関数上に構築された表現力豊かなオブジェクト指向ニューラルネットワーク:
import equinox as eqx
import jax
import jax.numpy as jnp
class SimpleMLP(eqx.Module):
linear1: eqx.nn.Linear
linear2: eqx.nn.Linear
def __init__(self, in_features, out_features, key):
k1, k2 = jax.random.split(key)
self.linear1 = eqx.nn.Linear(in_features, 64, key=k1)
self.linear2 = eqx.nn.Linear(64, out_features, key=k2)
def __call__(self, x):
x = jax.nn.relu(self.linear1(x))
return self.linear2(x)
key = jax.random.PRNGKey(0)
model = SimpleMLP(in_features=10, out_features=2, key=key)
sample_input = jnp.ones((10,))
output = model(sample_input)
print("Model output:", output)
共通パターン
vmap によるバッチベクトル化
問題: 手動の Python ループで関数をバッチ処理すると実行が遅くなります。
解決策:
jax.vmap を使用して単一サンプル関数を自動ベクトル化:
def predict_single(w, x):
return jnp.dot(x, w)
# バッチ次元（x の軸 0）に沿って自動ベクトル化
predict_batch = jax.vmap(predict_single, in_axes=(None, 0))
batch_x = jnp.ones((100, 10))
weights = jnp.ones(10)
predictions = predict_batch(weights, batch_x) # 出力形状: (100,)
ベストプラクティス (2026)
- 推奨: 純粋関数型コードを記述する: 副作用ゼロ、インプレース配列変更なし（
x[idx] = valの代わりにx.at[idx].set(val)を使用）。 - 推奨: パフォーマンスが重要な関数には
@jax.jitを付けて XLA コンパイルをトリガーする。 - 推奨: 乱数を生成するたびに PRNG キーを明示的に分割する (
jax.random.split(key))。 - 推奨: 生のパラメータ辞書を使う代わりに、Equinox や Flax Linen などの現代的な高レベルライブラリを使用する。
- 非推奨: JIT 関数内の動的トレーサーで標準の Python 条件文 (
if x > 0:) を使用する; 代わりにjax.lax.condを使用する。 - 非推奨: PRNG キーを再利用する; 再利用すると統計的に相関した乱数が生成される。
- 非推奨:
jit、grad、vmapで変換された関数内でグローバル変数をミューテーションする。
トラブルシューティング
| エラー | 原因 | 解決策 |
|---|---|---|
ConcretizationTypeError: Abstract tracer value used in Python if/while | Python の条件分岐が jit 内の動的 JAX 配列値に依存している。 | 標準の Python if の代わりに jax.lax.cond または jax.lax.while_loop を使用する。 |
JAX arrays are immutable | インプレース代入が試行された (arr[0] = 5)。 | 関数的更新構文を使用: arr = arr.at[0].set(5)。 |
PRNGKey reuse warning / identical randomness | 同一の PRNG キーを再利用すると同一の乱数が生成される。 | 各乱数操作の前にキーを分割: key, subkey = jax.random.split(key)。 |