JAX核心設計解析：函數式編程讓代碼更可控

很多人剛接觸JAX都會有點懵——參數為啥要單獨傳？隨機數還要自己管key？這跟PyTorch的畫風完全不一樣啊。

其實根本原因就一個：JAX是函數式編程而不是面向對象那套，想明白這點很多設計就都說得通了。

先說個核心區別

PyTorch里，模型是個對象，權重藏在里面，訓練的時候自己更新自己。這是典型的面向對象思路，狀態封裝在對象內部。

JAX的思路完全反過來。模型定義是模型定義，參數是參數，兩邊分得清清楚楚。函數本身不持有任何狀態，每次調用都把參數從外面傳進去。

這么做的好處？JAX可以把你的函數當純數學表達式來處理。求導、編譯、并行，想怎么折騰都行，因為函數里沒有藏著掖著的東西，行為完全可預測。

代碼對比一下就明白了

PyTorch這么寫：

import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 1)
def forward(self, x):
return self.linear(x)
model = Model()
x = torch.randn(5, 10)
output = model(x)

權重在self.linear里，模型自己管自己。

JAX配Flax是這樣：

import jax
import jax.numpy as jnp
from flax import linen as nn
class Model(nn.Module):
@nn.compact
def __call__(self, x):
return nn.Dense(1)(x)
model = Model()
key = jax.random.PRNGKey(0)
dummy = jnp.ones((1, 10))
params = model.init(key, dummy)['params']
x = jnp.ones((5, 10))
output = model.apply({'params': params}, x)

參數要先init出來，用的時候再apply進去。麻煩是麻煩了點，但參數流向一目了然，想做什么騷操作都很方便。

隨機數那個key是怎么回事

這個確實是JAX最讓新手頭疼的地方。不能直接random.normal()完事，非得帶個key：

key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (3,))

原因還是那個——函數式編程不允許隱藏狀態。

普通框架的隨機數生成器內部維護一個種子狀態，每次調用偷偷改一下。JAX不干這事。你得顯式給它一個key，它用完就扔，下次想生成隨機數再給個新的。

好處是隨機性完全可控可復現。jit編譯、多卡訓練、梯度計算，不管代碼怎么變換，只要key一樣結果就一樣。調試的時候不會遇到那種"明明代碼沒改怎么結果不一樣了"的玄學問題。

key不能復用，用之前要split

還有個規矩：同一個key只能用一次。要生成多個隨機數，得先split：

key = jax.random.PRNGKey(0)
key, subkey = jax.random.split(key)
a = jax.random.normal(subkey)
key, subkey = jax.random.split(key)
b = jax.random.uniform(subkey)

每次split出來的subkey都是獨立的隨機源。這套機制在分布式場景下特別香，不同機器拿不同的key，隨機性既獨立又可追溯。

合在一起看個完整例子

def forward(params, x):
w, b = params
return w * x + b
def init_params(key):
key_w, key_b = jax.random.split(key)
w = jax.random.normal(key_w)
b = jax.random.normal(key_b)
return w, b
key = jax.random.PRNGKey(0)
params = init_params(key)
x = jnp.array(2.0)
output = forward(params, x)

forward是純函數，輸入決定輸出，沒有副作用。隨機性在init_params里一次性處理完。參數獨立存放，想存哪存哪。

這種代碼JAX處理起來特別順手——jit編譯、自動微分、vmap批處理、多卡并行，都是開箱即用。

什么場景下JAX更合適

說實話JAX學習曲線是陡了點。但有些場景下它的優勢很明顯：做研究需要魔改模型結構的時候；物理仿真對數值精度和可復現性要求高的時候；大規模分布式訓練不想被隱藏狀態坑的時候；想自己擼optimizer或者自定義layer的時候。

適應了這套顯式風格之后其實挺舒服的。參數在哪、隨機數哪來的、函數干了啥，全都擺在明面上。沒有黑魔法，debug的時候心里有底。

https://avoid.overfit.cn/post/52fcdfd1d8054dcbb31783ed0547850e

作者：Ali Nawaz