Update translation: lectures/numpy_vs_numba_vs_jax.md

Author: mmcky

lectures/numpy_vs_numba_vs_jax.md

@@ -17,16 +17,15 @@ translation:
     Vectorized operations::Parallelized Numba: 并行化的 Numba
     Vectorized operations::Vectorized code with JAX: 使用 JAX 的向量化代码
     Vectorized operations::JAX plus vmap: JAX 加 vmap
-    Vectorized operations::JAX plus vmap::Version 1: 版本 1
-    Vectorized operations::vmap version 2: vmap 版本 2
     Vectorized operations::Summary: 总结
     Sequential operations: 顺序运算
     Sequential operations::Numba Version: Numba 版本
     Sequential operations::JAX Version: JAX 版本
     Sequential operations::Summary: 总结
+    Overall recommendations: 总体建议
 ---
 
-(parallel)=
+(numpy_numba_jax)=
 ```{raw} jupyter
 <div id="qe-notebook-header" align="right" style="text-align:right;">
         <a href="https://quantecon.org/" title="quantecon.org">
@@ -144,14 +143,13 @@ for x in grid:
             m = z
 ```
 
-
 ### NumPy 向量化
 
 如果我们切换到 NumPy 风格的向量化，就可以使用更大的网格，并且代码执行速度相对较快。
 
 这里我们使用 `np.meshgrid` 来创建二维输入网格 `x` 和 `y`，使得 `f(x, y)` 能生成乘积网格上的所有计算结果。
 
-（这一策略可以追溯到 Matlab。）
+（这一策略可以追溯到 MATLAB。）
 
 ```{code-cell} ipython3
 grid = np.linspace(-3, 3, 3_000)
@@ -169,7 +167,6 @@ print(f"NumPy result: {z_max_numpy:.6f}")
 
 （并行化效率不高，因为二进制文件在看到数组 `x` 和 `y` 的大小之前就已经被编译了。）
 
-
 ### 与 Numba 的比较
 
 现在让我们看看能否使用简单循环的 Numba 获得更好的性能。
@@ -208,7 +205,6 @@ with qe.Timer(precision=8):
 
 另一方面，Numba 例程使用的内存少得多，因为我们只处理一个一维网格。
 
-
 ### 并行化的 Numba
 
 现在让我们使用 `prange` 尝试 Numba 的并行化：
@@ -231,24 +227,44 @@ def compute_max_numba_parallel(grid):
 
 ```
 
-这通常会返回不正确的结果：
+这将返回 `-inf`——即 `m` 的初始值，仿佛它从未被更新过：
 
 ```{code-cell} ipython3
 z_max_parallel_incorrect = compute_max_numba_parallel(grid)
 print(f"Numba result: {z_max_parallel_incorrect} 😱")
 ```
 
-原因是变量 `m` 被多个线程共享，但没有得到正确控制。
+要理解原因，请回忆 `prange` 会将外层循环拆分到各个线程中。
+
+每个线程都会得到自己的 `m` 私有副本，初始化为 `-np.inf`，并在其负责的迭代块中正确地更新它。
 
-当多个线程同时尝试读写 `m` 时，它们会相互干扰。
+但在循环结束时，Numba 需要将各线程的 `m` 副本合并为一个单一的值——即**归约**操作。
 
-线程读取了 `m` 的过时值，或者相互覆盖了更新——或者 `m` 始终保持其初始值而从未被更新。
+对于它能识别的模式，例如 `m += z`（求和）或 `m = max(m, z)`（求最大值），Numba 知道合并算子。
 
-这里有一个更仔细编写的版本。
+但它无法将 `if z > m: m = z` 识别为最大值归约，因此各线程的结果永远不会被合并，`m` 始终保持其初始值。
+
+最简单的修复方法是将条件判断替换为 Numba 能识别的 `max`：
 
 ```{code-cell} ipython3
 @numba.jit(parallel=True)
 def compute_max_numba_parallel(grid):
+    n = len(grid)
+    m = -np.inf
+    for i in numba.prange(n):
+        for j in range(n):
+            x = grid[i]
+            y = grid[j]
+            z = np.cos(x**2 + y**2) / (1 + x**2 + y**2)
+            m = max(m, z)
+    return m
+```
+
+另一种方法是使循环体在不同 `i` 之间完全独立，并自行处理归约：
+
+```{code-cell} ipython3
+@numba.jit(parallel=True)
+def compute_max_numba_parallel_v2(grid):
     n = len(grid)
     row_maxes = np.empty(n)
     for i in numba.prange(n):
@@ -263,9 +279,7 @@ def compute_max_numba_parallel(grid):
     return np.max(row_maxes)
 ```
 
-现在 `for i in numba.prange(n)` 所作用的代码块在不同的 `i` 之间是独立的。
-
-每个线程写入数组 `row_maxes` 的不同元素，并行化是安全的。
+在这里，每个线程写入 `row_maxes` 的不同元素，因此我们通过 `np.max` 自行处理归约。
 
 ```{code-cell} ipython3
 z_max_parallel = compute_max_numba_parallel(grid)
@@ -283,7 +297,6 @@ with qe.Timer(precision=8):
 
 对于更强大的机器和更大的网格尺寸，即使在 CPU 上，并行化也能带来显著的速度提升。
 
-
 ### 使用 JAX 的向量化代码
 
 表面上，JAX 中的向量化代码与 NumPy 代码类似。
@@ -304,7 +317,7 @@ def f(x, y):
 
 ```{code-cell} ipython3
 grid = jnp.linspace(-3, 3, 3_000)
-x_mesh, y_mesh = np.meshgrid(grid, grid)
+x_mesh, y_mesh = jnp.meshgrid(grid, grid)
 
 with qe.Timer(precision=8):
     z_max = jnp.max(f(x_mesh, y_mesh))
@@ -321,14 +334,13 @@ with qe.Timer(precision=8):
     z_max.block_until_ready()
 ```
 
-编译完成后，由于 GPU 加速，JAX 明显快于 NumPy。
+编译完成后，JAX 明显快于 NumPy，在 GPU 上尤为如此。
 
 编译开销是一次性成本，当函数被反复调用时，这种开销是值得的。
 
-
 ### JAX 加 vmap
 
-NumPy 代码和 JAX 代码都存在一个问题：
+NumPy 代码和上述 JAX 代码都存在一个问题：
 
 虽然扁平数组占用内存较少
 
@@ -346,9 +358,9 @@ x_mesh.nbytes + y_mesh.nbytes
 
 幸运的是，JAX 提供了一种使用 [jax.vmap](https://docs.jax.dev/en/latest/_autosummary/jax.vmap.html) 的不同方法。
 
-#### 版本 1
+`vmap` 的思路是将向量化分阶段进行，将一个对单个值进行操作的函数转化为对数组进行操作的函数。
 
-以下是我们应用 `vmap` 的一种方式。
+以下是我们将其应用于当前问题的方式。
 
 ```{code-cell} ipython3
 # 设置 f，使其在给定任意 y 时，对所有 x 计算 f(x, y)
@@ -375,32 +387,19 @@ with qe.Timer(precision=8):
     z_max.block_until_ready()
 ```
 
-通过避免使用大型输入数组 `x_mesh` 和 `y_mesh`，这个 `vmap` 版本使用的内存少得多。
-
-在 CPU 上运行时，其运行时间与网格版本相似。
-
-在 GPU 上运行时，通常速度要快得多。
-
-实际上，使用 `vmap` 还有另一个优势：它允许我们将向量化分阶段进行。
-
-这往往会产生比传统向量化代码更易于理解的代码。
+通过避免使用大型输入数组 `x_mesh` 和 `y_mesh`，这个 `vmap` 版本使用的内存少得多，运行时间也相近。
 
-当我们处理更大的问题时，将进一步探讨这些想法。
+但我们仍然留有一些速度提升的空间未被利用。
 
+上面的代码计算了完整的二维数组 `f(x,y)`，然后再取最大值。
 
-### vmap 版本 2
+此外，`jnp.max` 调用位于 JIT 编译函数 `f` 之外，因此编译器无法将这些操作融合为单个内核。
 
-我们可以使用 vmap 进一步提高内存效率。
-
-在前一个版本中，虽然我们避免了大型输入数组，但在计算最大值之前仍然会创建大型输出数组 `f(x,y)`。
-
-让我们尝试一种略有不同的方法，将求最大值操作移到内部。
-
-由于这一改变，我们永远不会计算二维数组 `f(x,y)`。
+我们可以通过将最大值操作移到内部并将所有内容包装在一个 `@jax.jit` 中来解决这两个问题：
 
 ```{code-cell} ipython3
 @jax.jit
-def compute_max_vmap_v2(grid):
+def compute_max_vmap(grid):
     # 构建一个沿每行取最大值的函数
     f_vec_x_max = lambda y: jnp.max(f(grid, y))
     # 向量化该函数，以便我们可以同时对所有行调用
@@ -416,25 +415,26 @@ def compute_max_vmap_v2(grid):
 
 我们将此函数应用于所有行，然后取各行最大值中的最大值。
 
+由于将最大值操作移到内部，我们永远不会构建完整的二维数组 `f(x,y)`，从而节省了更多内存。
+
+并且由于所有内容都在单个 `@jax.jit` 下，编译器可以将所有操作融合为一个优化的内核。
+
 让我们试试。
 
 ```{code-cell} ipython3
 with qe.Timer(precision=8):
-    z_max = compute_max_vmap_v2(grid).block_until_ready()
+    z_max = compute_max_vmap(grid).block_until_ready()
 
-print(f"JAX vmap v1 result: {z_max:.6f}")
+print(f"JAX vmap result: {z_max:.6f}")
 ```
 
 让我们再次运行以消除编译时间：
 
 ```{code-cell} ipython3
 with qe.Timer(precision=8):
-    z_max = compute_max_vmap_v2(grid).block_until_ready()
+    z_max = compute_max_vmap(grid).block_until_ready()
 ```
 
-如果您像我们一样在 GPU 上运行，应该能看到又一个不小的速度提升。
-
-
 ### 总结
 
 在我们看来，JAX 是向量化运算的赢家。
@@ -449,15 +449,13 @@ with qe.Timer(precision=8):
 
 对于经济学、计量经济学和金融学中遇到的大多数情况，将高效并行化的工作交给 JAX 编译器，远比尝试手工编写这些例程要好得多。
 
-
 ## 顺序运算
 
 某些运算本质上是顺序的——因此难以或不可能向量化。
 
 在这种情况下，NumPy 是一个较差的选择，我们只剩下 Numba 或 JAX 可以选择。
 
-为了比较这两种选择，我们将重新回顾在{doc}`Numba 讲座 <numba>`中看到的迭代二次映射问题。
-
+为了比较这两种选择，我们将重新回顾在 {doc}`Numba 讲座 <numba>` 中看到的迭代二次映射问题。
 
 ### Numba 版本
 
@@ -545,8 +543,7 @@ with qe.Timer(precision=8):
 
 JAX 对于这种顺序运算也相当高效。
 
-JAX 和 Numba 在编译后都能提供出色的性能，对于纯顺序运算，Numba 通常（但并非总是）提供略快的速度。
-
+JAX 和 Numba 在编译后都能提供出色的性能。
 
 ### 总结
 
@@ -560,4 +557,31 @@ Numba 版本简单直观，易于阅读：我们只需分配一个数组，然
 
 此外，JAX 的不可变数组意味着我们无法简单地就地更新数组元素，这使得直接复制 Numba 使用的算法变得困难。
 
-对于这类顺序运算，在代码清晰度、实现便利性以及高性能方面，Numba 是明显的赢家。
+对于这类顺序运算，在代码清晰度和实现便利性方面，Numba 是明显的赢家。
+
+## 总体建议
+
+让我们退一步，总结一下各方的权衡取舍。
+
+对于**向量化操作**，JAX 是最强的选择。
+
+得益于 JIT 编译和在 CPU 与 GPU 上的高效并行化，它在速度上与 NumPy 持平或超越 NumPy。
+
+`vmap` 变换降低了内存使用量，并且通常比传统的基于网格的向量化产生更清晰的代码。
+
+此外，JAX 函数支持自动微分，我们将在 {doc}`autodiff` 中进一步探讨。
+
+对于**顺序操作**，Numba 具有明显优势。
+
+代码自然且可读——只需一个带有装饰器的 Python 循环——性能也非常出色。
+
+JAX 可以通过 `lax.scan` 处理顺序问题，但语法不够直观。
+
+```{note}
+`lax.scan` 的一个重要优势是它支持通过循环进行自动微分，而 Numba 无法做到这一点。
+如果您需要对顺序计算进行微分（例如，计算轨迹对模型参数的敏感性），尽管语法不够自然，JAX 仍是更好的选择。
+```
+
+在实践中，许多问题涉及两种模式的混合。
+
+一个实用的经验法则是：对于新项目默认使用 JAX，尤其是当硬件加速或可微分性可能有用时，而当您有一个需要快速且可读的紧凑顺序循环时，则选用 Numba。