JAX
JAX 是「NumPy 強化版」。它結合了 Autograd（自動微分）與 XLA（編譯）。2025 年 Flax NNX（PyTorch 風格的物件導向程式設計）已成為標準。
使用時機
- 高效能數值與科學運算：利用 Autograd 與 XLA 編譯在 GPU/TPU 上加速線性代數運算。
- 客製深度學習研究：使用 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("梯度:", 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 能自動處理二維矩陣，無需迴圈
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("計算出的梯度:", 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("常態分佈樣本平均數:", 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("模型輸出:", 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.at[idx].set(val)而非x[idx] = 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 陣列值。 | 使用 jax.lax.cond 或 jax.lax.while_loop 取代標準 Python if。 |
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)。 |