对比PyTorch、TensorFlow、JAX、Theano，我发现都在关注两大问题

作者｜王益

OneFlow社区编译

翻译｜杨婷

最近，我在处理 PyTorch 分布式和 TorchRec 相关的工作，为此，我开始学习 PyTorch 2.0。在业余时间，我也在跟着Alpa作者学习JAX和XLA。如今回顾这些技术，我发现它们的关注点似乎都是如下两个问题：

1. 包含自动求导和并行在内的函数转换，例如 vmap, pmap 和 pjit 等；

2. 异构计算，CPU 负责控制流，GPU/TPU 负责张量计算和集合通信。

本文档中的所有例子都支持在 Colab 中运行：

1

函数转换

“函数转换”意为将一个程序转变成另一个程序，最常见的例子是自动求导（autograd）。自动求导采用用户编写的前向过程并创建后向过程，对于用户来说，编写自动求导通常都太过复杂。函数转换的主要难点在于：在编写函数转换算法时以何种方式表示输入和输出过程。

Theano：显式地构建 IR

Theano是最早的深度学习工具之一，也就是如今为人们所熟知的Aesara项目。Theano有一个允许用户在内存中将IR构建为数据结构的API，因此Theano可实现自动求导，并将结果输出为 Python 函数。
 

import aesara
from aesara import tensor as at
 
 
a = at.dscalar("a") # Define placeholders, which have no values.
b = at.dscalar("b")
 
 
c = a * b              # c now contains the IR of an expression.TT
dc = aesara.grad(c, a) # Convert the IR in c into another one, dc
 
 
f_dc = aesara.function([a, b], dc) # Convert the IR into a Python function,
assert f_dc(1.5, 2.5) == 2.5       # so we can call it.

TensorFlow 1.x：用于运行 IR 的虚拟机

TensorFlow 1.x明确保留了构建IR的想法。若在TensorFlow中运行上述示例，结果不会有什么差别；但倘若在TensorFlow 1.x中来运行，最大的差别在