第1章 JAX与PyTorch生态全景对比

一、JAX与PyTorch生态全景对比：为什么需要互转？

说实话，我入行深度学习那会儿，PyTorch还没出生呢。那时候大家用TensorFlow 1.x，写个模型跟受刑似的。后来PyTorch横空出世，动态图让调试变得无比丝滑。再后来，JAX带着函数式编程和XLA编译杀了出来。

现在很多团队两个框架都在用。我见过不少项目，研究阶段用PyTorch快速验证，生产部署时转成JAX榨取性能。所以，搞懂它们为什么需要互转，是咱们这门课的第一块基石。

1.1 核心差异：动态图 vs 静态图

这是最根本的区别。我刚开始接触时也绕了很久。

特性	PyTorch（动态图）	JAX（静态图）
执行方式	边定义边执行	先编译再执行
调试体验	直接打断点，print大法好	编译后像黑盒，调试靠技巧
性能优化	依赖Python运行时	XLA编译，极致优化
控制流	原生Python if/for	需要特殊处理（如lax.cond）

动态图就像你写一行代码，计算机立刻执行一行。你随时可以print看看中间结果。我当年调试一个Transformer模型，就是靠逐层打印tensor形状定位到维度不对的。这种即时反馈，对研究和开发太友好了。

静态图则不同。你先搭好整个计算图，然后一次性编译执行。JAX的jit装饰器就是这么干的。好处是编译器能看到全局，可以做算子融合、内存复用等优化。坏处是——嗯，你没法在编译后的函数里随便print。

我的经验：研究阶段用PyTorch快速迭代，生产阶段用JAX榨干硬件性能。两者不是替代关系，而是互补关系。

1.2 函数式 vs 面向对象

这个差异，说实话，一开始让我很不适应。

PyTorch是典型的面向对象风格。你定义一个nn.Module类，里面挂着self.weight、self.bias这些参数。训练时调用model.parameters()获取参数列表，传给优化器。一切都很自然，对吧？

# PyTorch 风格
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 2)
    
    def forward(self, x):
        return self.fc(x)

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())

JAX则是函数式的。参数是显式传入函数的，模型本身没有内部状态。我刚开始写JAX时，总想找"模型对象"在哪，后来才明白——JAX里模型就是一堆数组，函数就是纯函数。

# JAX 风格
def forward(params, x):
    w, b = params
    return jnp.dot(x, w) + b

params = (init_w, init_b)  # 参数是显式的
loss = jit(lambda p, x, y: jnp.mean((forward(p, x) - y) ** 2))

避坑指南：我曾经把一个PyTorch模型直接照搬到JAX，结果发现参数更新逻辑完全不一样。PyTorch的优化器会帮你管理状态，JAX需要你手动处理参数更新。这个坑，咱们后面章节会详细讲。

1.3 适用场景分析

那到底什么时候用PyTorch，什么时候用JAX？我根据自己的项目经验，整理了一张表：

场景	推荐框架	原因
快速原型开发	PyTorch	动态图调试方便，社区生态成熟
学术研究	PyTorch	论文代码大多用PyTorch，复现容易
大规模分布式训练	JAX	pmap/shard_map天然支持多设备
生产部署（低延迟）	JAX	XLA编译后推理速度极快
强化学习	两者皆可	PyTorch生态更全，JAX性能更好
自定义算子/梯度	JAX	vmap/jacfwd/jacrev等自动微分工具强大

说白了，如果你追求开发效率和生态丰富度，PyTorch是首选。如果你追求极致性能和可扩展性，JAX更合适。但现实是——很多项目需要两者兼顾。

1.4 为什么需要互转？

我遇到过最典型的场景是这样的：

研究团队用PyTorch开发了一个新模型，发了论文
工程团队需要把这个模型部署到生产环境，要求低延迟
PyTorch的推理性能达不到要求，需要转成JAX用XLA编译
于是，模型互转就成了刚需

还有另一种情况：你看到一个很棒的JAX项目（比如DeepMind的AlphaFold），但你的团队只熟悉PyTorch。这时候，把JAX模型转成PyTorch，就能复用团队已有的训练pipeline。

注意：互转不是简单的"翻译代码"。PyTorch和JAX在随机数生成、参数初始化、BatchNorm等细节上都有差异。直接硬转，模型精度可能会掉。咱们这门课会逐一攻克这些坑。

1.5 知识体系总览

下面这张图，是我梳理的JAX与PyTorch互转的核心知识体系。你可以把它当作整个课程的地图：

你看，互转不是简单的"翻译代码"。它涉及张量转换、模型结构映射、训练流程适配、推理部署优化等多个层面。每个层面都有各自的坑和技巧。

我的建议：别想着一次性搞定所有互转。先从最简单的张量转换开始，逐步深入到模型结构、训练流程。咱们这门课就是按这个思路设计的。

好了，第一章就到这里。咱们把JAX和PyTorch的差异、各自的适用场景、以及为什么需要互转都理清楚了。下一章，我会带大家动手做张量互转——这是所有互转操作的基础。

📚 课程目录

一、JAX与PyTorch生态全景对比：为什么需要互转？

1.1 核心差异：动态图 vs 静态图

1.2 函数式 vs 面向对象

1.3 适用场景分析

1.4 为什么需要互转？

1.5 知识体系总览