Update translation: lectures/jax_intro.md

Author: mmcky

lectures/jax_intro.md

@@ -21,16 +21,15 @@ translation:
     JAX as a NumPy Replacement::Differences::Size Experiment: 大小实验
     JAX as a NumPy Replacement::Differences::Precision: 精度
     JAX as a NumPy Replacement::Differences::Immutability: 不可变性
-    JAX as a NumPy Replacement::Differences::A workaround: 变通方法
+    JAX as a NumPy Replacement::Differences::A Workaround: 变通方法
     Functional Programming: 函数式编程
     Functional Programming::Pure functions: 纯函数
-    Functional Programming::Examples: 示例
-    Functional Programming::Why Functional Programming?: 为什么使用函数式编程？
+    Functional Programming::Examples -- Pure and Impure: 示例——纯函数与非纯函数
+    Functional Programming::Why Functional Programming?: 为什么要函数式编程？
     Random numbers: 随机数
-    Random numbers::Random number generation: 随机数生成
-    Random numbers::Why explicit random state?: 为什么要显式随机状态？
-    Random numbers::Why explicit random state?::NumPy's approach: NumPy 的方法
-    Random numbers::Why explicit random state?::JAX's approach: JAX 的方法
+    Random numbers::NumPy / MATLAB Approach: NumPy / MATLAB 方法
+    Random numbers::JAX: JAX
+    Random numbers::Benefits: 优势
     JIT Compilation: JIT 编译
     JIT Compilation::With NumPy: 使用 NumPy
     JIT Compilation::With JAX: 使用 JAX
@@ -343,19 +342,20 @@ a
 * 不会改变全局状态
 * 不会修改传递给函数的数据（不可变数据）
 
-### 示例
+### 示例——纯函数与非纯函数
 
 以下是一个*非纯*函数的示例：
 
 ```{code-cell} ipython3
 tax_rate = 0.1
-prices = [10.0, 20.0]
 
 def add_tax(prices):
     for i, price in enumerate(prices):
         prices[i] = price * (1 + tax_rate)
-    print('Post-tax prices: ', prices)
-    return prices
+
+prices = [10.0, 20.0]
+add_tax(prices)
+prices
 ```
 
 这个函数不是纯函数，因为：
@@ -366,15 +366,21 @@ def add_tax(prices):
 以下是一个*纯*版本：
 
 ```{code-cell} ipython3
-tax_rate = 0.1
-prices = (10.0, 20.0)
 
 def add_tax_pure(prices, tax_rate):
     new_prices = [price * (1 + tax_rate) for price in prices]
     return new_prices
+
+tax_rate = 0.1
+prices = (10.0, 20.0)
+after_tax_prices = add_tax_pure(prices, tax_rate)
+after_tax_prices
 ```
 
-这个纯版本通过函数参数使所有依赖关系变得明确，并且不修改任何外部状态。
+这是纯函数，因为：
+
+* 所有依赖关系通过函数参数显式表达
+* 不修改任何外部状态
 
 ### 为什么要函数式编程？
 
@@ -404,13 +410,29 @@ JAX 使用函数式编程风格，以便用户构建的函数能够直接映射
 
 JAX 中的随机数生成与 NumPy 或 MATLAB 中的模式有很大不同。
 
-起初，您可能会觉得语法相当冗长。
+### NumPy / MATLAB 方法
 
-但为了维护我们刚刚讨论的函数式编程风格，这种语法和语义是必要的。
+在 NumPy / MATLAB 中，生成通过维护隐藏的全局状态来工作。
 
-此外，对随机状态的完全控制对于并行编程至关重要，例如当我们想要沿多个线程运行独立实验时。
+```{code-cell} ipython3
+np.random.seed(42)
+print(np.random.randn(2))   
+```
 
-### 随机数生成
+每次我们调用随机函数时，隐藏状态都会被更新：
+
+```{code-cell} ipython3
+print(np.random.randn(2)) 
+```
+
+这个函数*不是纯函数*，因为：
+
+* 它是非确定性的：相同的输入，不同的输出
+* 它有副作用：它修改了全局随机数生成器状态
+
+这在并行化下是危险的——必须仔细控制每个线程中发生的事情。
+
+### JAX
 
 在 JAX 中，随机数生成器的状态被显式控制。
 
@@ -525,125 +547,58 @@ plt.show()
 下面的函数使用 `split` 生成 `k` 个（准）独立的随机 `n x n` 矩阵。
 
 ```{code-cell} ipython3
-def gen_random_matrices(key, n=2, k=3):
+def gen_random_matrices(
+        key,   # JAX key for random numbers
+        n=2,   # Matrices will be n x n
+        k=3    # Number of matrices to generate
+    ):
     matrices = []
     for _ in range(k):
         key, subkey = jax.random.split(key)
         A = jax.random.uniform(subkey, (n, n))
         matrices.append(A)
-        print(A)
     return matrices
 ```
 
 ```{code-cell} ipython3
 seed = 42
 key = jax.random.key(seed)
-matrices = gen_random_matrices(key)
-```
-
-我们也可以在循环迭代时使用 `fold_in`：
-
-```{code-cell} ipython3
-def gen_random_matrices(key, n=2, k=3):
-    matrices = []
-    for i in range(k):
-        step_key = jax.random.fold_in(key, i)
-        A = jax.random.uniform(step_key, (n, n))
-        matrices.append(A)
-        print(A)
-    return matrices
+gen_random_matrices(key)
 ```
 
-```{code-cell} ipython3
-key = jax.random.key(seed)
-matrices = gen_random_matrices(key)
-```
-
-### 为什么要显式随机状态？
-
-为什么 JAX 需要这种相对冗长的随机数生成方法？
-
-一个原因是为了维护纯函数。
-
-让我们通过比较 NumPy 和 JAX 来看看随机数生成与纯函数的关系。
-
-#### NumPy 的方法
-
-在 NumPy 的旧版随机数生成 API（模仿 MATLAB）中，生成通过维护隐藏的全局状态来工作。
-
-每次我们调用随机函数时，这个状态都会被更新：
-
-```{code-cell} ipython3
-np.random.seed(42)
-print(np.random.randn())   # Updates state of random number generator
-print(np.random.randn())   # Updates state of random number generator
-```
-
-每次调用都返回不同的值，即使我们用相同的输入（没有参数）调用相同的函数。
-
-这个函数*不是纯函数*，因为：
-
-* 它是非确定性的：相同的输入（在这种情况下，没有输入）产生不同的输出
-* 它有副作用：它修改了全局随机数生成器状态
-
-#### JAX 的方法
-
-如上所示，JAX 采用了不同的方法，通过密钥使随机性显式化。
-
-例如：
-
-```{code-cell} ipython3
-def random_sum_jax(key):
-    key1, key2 = jax.random.split(key)
-    x = jax.random.normal(key1)
-    y = jax.random.normal(key2)
-    return x + y
-```
-
-使用相同的密钥，我们总是得到相同的结果：
-
-```{code-cell} ipython3
-key = jax.random.key(42)
-random_sum_jax(key)
-```
+这个函数是*纯函数*
 
-```{code-cell} ipython3
-random_sum_jax(key)
-```
-
-要获得新的抽取，我们需要提供一个新密钥。
-
-函数 `random_sum_jax` 是纯函数，因为：
-
-* 它是确定性的：相同的密钥总是产生相同的输出
+* 确定性：相同的输入，相同的输出
 * 无副作用：没有隐藏状态被修改
 
-JAX 的显式性带来了显著的好处：
+### 优势
+
+如上所述，这种显式性是有价值的：
 
 * 可复现性：通过重用密钥轻松重现结果
-* 并行化：每个线程可以拥有自己的密钥而不会产生冲突
-* 调试：没有隐藏状态使代码更容易推理
+* 并行化：控制每个独立线程中发生的事情
+* 调试：没有隐藏状态使代码更容易测试
 * JIT 兼容性：编译器可以更积极地优化纯函数
 
-最后一点将在下一节中进行扩展。
-
 ## JIT 编译
 
 JAX 的即时（JIT）编译器通过生成随任务大小和硬件变化的高效机器码来加速执行。
 
 我们在 {ref}`上文 <jax_speed>` 中已经看到了 JAX 的 JIT 编译器结合并行硬件的强大之处，当时我们对一个大数组应用了 `cos` 函数。
 
-让我们用一个更复杂的函数尝试同样的操作：
+这里我们研究针对更复杂函数的 JIT 编译。
+
+### 使用 NumPy
+
+我们先用 NumPy 试试：
 
 ```{code-cell}
 def f(x):
     y = np.cos(2 * x**2) + np.sqrt(np.abs(x)) + 2 * np.sin(x**4) - x**2
     return y
 ```
 
-### 使用 NumPy
-
-我们先用 NumPy 试试：
+让我们用较大的 `x` 运行：
 
 ```{code-cell}
 n = 50_000_000
@@ -656,9 +611,17 @@ with qe.Timer():
     y = f(x)
 ```
 
-### 使用 JAX
+**即时执行**模型
+
+* 每个操作在遇到时立即执行，在下一个操作开始之前将其结果实体化。
 
-现在让我们用 JAX 再试一次。
+缺点
+
+* 并行化程度极低
+* 内存占用大——产生许多中间数组
+* 大量内存读写
+
+### 使用 JAX
 
 作为第一步，我们将整个代码中的 `np` 替换为 `jnp`：
 
@@ -691,11 +654,20 @@ with qe.Timer():
 
 结果与 `cos` 示例类似——JAX 更快，尤其是在 JIT 编译后的第二次运行中。
 
-然而，使用 JAX，我们还有另一个技巧——我们可以对整个函数进行 JIT 编译，而不仅仅是单个操作。
+这是因为单个数组操作在 GPU 上并行化了。
+
+但我们仍然在使用即时执行模式
+
+* 由于中间数组导致大量内存占用
+* 大量内存读写
+
+此外，在 GPU 上还启动了许多独立的内核。
 
 ### 编译整个函数
 
-JAX 即时（JIT）编译器可以通过将数组运算融合到单个优化内核中来加速函数内部的执行。
+幸运的是，使用 JAX，我们还有另一个技巧——我们可以对整个函数进行 JIT 编译，而不仅仅是单个操作。
+
+编译器将所有数组操作融合到单个优化内核中。
 
 让我们用函数 `f` 来试试这个：
 
@@ -719,9 +691,12 @@ with qe.Timer():
     jax.block_until_ready(y);
 ```
 
-运行时间再次改善——现在是因为我们融合了所有操作，使编译器能够更积极地进行优化。
+运行时间再次改善——现在是因为我们融合了所有操作
 
-例如，编译器可以消除对硬件加速器的多次调用以及许多中间数组的创建。
+* 基于整个计算序列的激进优化
+* 消除对硬件加速器的多次调用
+
+内存占用也大幅降低——不再创建中间数组。
 
 顺便提一下，当针对 JIT 编译器的函数时，更常见的语法是：
 
@@ -741,11 +716,9 @@ XLA 随后将这些操作融合并优化为针对可用硬件（CPU、GPU 或 TP
 
 ### 编译非纯函数
 
-现在我们已经看到了 JIT 编译的强大之处，理解它与纯函数的关系非常重要。
-
-虽然 JAX 在编译非纯函数时通常不会抛出错误，但执行会变得不可预测。
+虽然 JAX 在编译非纯函数时通常不会抛出错误，但执行会变得不可预测！
 
-以下是一个使用全局变量的例子：
+以下是一个示例：
 
 ```{code-cell} ipython3
 a = 1  # global
@@ -787,65 +760,6 @@ f(x)
 
 这个故事的寓意：使用 JAX 时请编写纯函数！
 
-## 使用 `vmap` 进行向量化
-
-JAX 的另一个强大变换是 `jax.vmap`，它能自动将一个针对单个输入编写的函数向量化，使其可以在批量数据上运行。
-
-这避免了手动编写向量化代码或使用显式循环的需要。
-
-### 一个简单的示例
-
-假设我们有一个函数，用于计算一组数字的均值与中位数之差。
-
-```{code-cell} ipython3
-def mm_diff(x):
-    return jnp.mean(x) - jnp.median(x)
-```
-
-我们可以将其应用于单个向量：
-
-```{code-cell} ipython3
-x = jnp.array([1.0, 2.0, 5.0])
-mm_diff(x)
-```
-
-现在假设我们有一个矩阵，想要对每一行计算这些统计量。
-
-不使用 `vmap` 时，我们需要显式循环：
-
-```{code-cell} ipython3
-X = jnp.array([[1.0, 2.0, 5.0],
-               [4.0, 5.0, 6.0],
-               [1.0, 8.0, 9.0]])
-
-for row in X:
-    print(mm_diff(row))
-```
-
-然而，Python 循环速度较慢，无法被 JAX 高效编译或并行化。
-
-使用 `vmap` 可以将计算保留在加速器上，并与其他 JAX 变换（如 `jit` 和 `grad`）组合使用：
-
-```{code-cell} ipython3
-batch_mm_diff = jax.vmap(mm_diff)
-batch_mm_diff(X)
-```
-
-函数 `mm_diff` 是针对单个数组编写的，而 `vmap` 自动将其提升为按行作用于矩阵的函数——无需循环，无需重新塑形。
-
-### 组合变换
-
-JAX 的优势之一在于各变换可以自然地组合使用。
-
-例如，我们可以对向量化函数进行 JIT 编译：
-
-```{code-cell} ipython3
-fast_batch_mm_diff = jax.jit(jax.vmap(mm_diff))
-fast_batch_mm_diff(X)
-```
-
-`jit`、`vmap` 以及（我们接下来将看到的）`grad` 的这种组合方式是 JAX 设计的核心，使其在科学计算和机器学习领域尤为强大。
-
 ## 练习