点击上方“Deephub Imba”,关注公众号,好文章不错过 !
很多人刚接触JAX都会有点懵——参数为啥要单独传？随机数还要自己管key？这跟PyTorch的画风完全不一样啊。
其实根本原因就一个：JAX是函数式编程而不是面向对象那套，想明白这点很多设计就都说得通了。
先说个核心区别
PyTorch里，模型是个对象，权重藏在里面，训练的时候自己更新自己。这是典型的面向对象思路，状态封装在对象内部。
JAX的思路完全反过来。模型定义是模型定义，参数是参数，两边分得清清楚楚。函数本身不持有任何状态，每次调用都把参数从外面传进去。
这么做的好处？JAX可以把你的函数当纯数学表达式来处理。求导、编译、并行，想怎么折腾都行，因为函数里没有藏着掖着的东西，行为完全可预测。
代码对比一下就明白了
PyTorch这么写：
importtorch
importtorch.nnasnn
classModel(nn.Module):
def__init__(self):
super().__init__()
self.linear=nn.Linear(10, 1)
defforward(self, x):
returnself.linear(x)
model=Model()
x=torch.randn(5, 10)
output=model(x)
权重在self.linear里，模型自己管自己。
JAX配Flax是这样：
importjax
importjax.numpyasjnp
fromflaximportlinenasnn
classModel(nn.Module):
@nn.compact
def__call__(self, x):
returnnn.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，随机性既独立又可追溯。
合在一起看个完整例子
defforward(params, x):
w, b=params
returnw*x+b
definit_params(key):
key_w, key_b=jax.random.split(key)
w=jax.random.normal(key_w)
b=jax.random.normal(key_b)
returnw, 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
喜欢就关注一下吧：
点个 在看 你最好看！