从PyTorch切换到JAX，真的不是换个键盘那么简单

把项目从 PyTorch 迁移到 JAX 之后，我才意识到这根本不是 API 的替换，而是整个编程范式的推倒重来。最让我崩溃的是在实现一个简单的自定义 Loss 函数时，陷入了没完没了的 TypeError 和数值异常中。

最典型的坑就是 JAX 的不可变性（Immutability）。在 PyTorch 里，我习惯直接对 Tensor 进行原位操作，比如 x += 1 或者 x[0] = 5。结果在 JAX 里这么写直接报错，因为它不允许修改数组。

报错信息大概长这样：

TypeError: 'DeviceArray' object does not support item assignment

排查之后才发现，JAX 要求必须使用 .at[].set() 这种函数式写法。比如我想更新某个索引的值，得写成：

# PyTorch 写法: x[idx] = val
# JAX 写法:
x = x.at[idx].set(val)

这种逻辑转换在小模型里还好，一旦涉及到复杂的 RNN 状态更新或者自定义的优化器，代码量直接翻倍，而且极容易因为漏写一个赋值号导致梯度完全没更新，模型原地打转。

另一个深坑是随机数种子（PRNGKey）。PyTorch 的 torch.manual_seed(42) 是全局状态，随处调用 torch.randn 都能拿到结果。但 JAX 强制要求显式传递 Key。我起初为了省事，在循环里一直用同一个 key，结果发现每次生成的随机数竟然一模一样，导致模型根本不收敛。

正确的排查路径是必须学会 jax.random.split。每次用完一个 key，必须把它“劈”开成两个，一个留着，一个传给下一步：

key, subkey = jax.random.split(key)
x = jax.random.normal(subkey, (10,))

最让我觉得心累的是 jit 编译。虽然 jax.jit 带来的加速很猛，但它对 Python 原生控制流（如 if/while）极其不友好。只要在 jit 函数里写了依赖于输入值的 if 判断，就会报 ConcretizationTypeError。我得把所有的 Python if 全改成 jax.lax.cond 或者 jax.lax.select，这让代码的可读性瞬间掉到了冰点。

总结这次迁移，PyTorch 像是在写面向对象，而 JAX 是在写纯函数式数学。如果你习惯了命令式编程，强行切换 JAX 可能会在头两周觉得在跟编译器打架。