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的时候心里有底。

作者：Ali Nawaz