Merge pull request #49 from QuantEcon/translation-sync-2026-04-12T13-13-55-pr-525

🌐 [translation-sync] Improve NumPy vs Numba vs JAX lecture

Author: HumphreyYang

.translate/state/numpy_vs_numba_vs_jax.md.yml

@@ -1,5 +1,5 @@
-source-sha: 05ce95691fd97e48da39dd6d58fe032c03e8813d
-synced-at: "2026-04-09"
+source-sha: 94dd7d22385ec46d740db1fc2cddf05c29377594
+synced-at: "2026-04-12"
 model: claude-sonnet-4-6
 mode: UPDATE
 section-count: 3

lectures/numpy_vs_numba_vs_jax.md

@@ -17,8 +17,6 @@ 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 版本
@@ -27,7 +25,7 @@ translation:
     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">
@@ -155,7 +153,7 @@ for x in grid:
 
 这里我们使用 `np.meshgrid` 来创建二维输入网格 `x` 和 `y`，使得 `f(x, y)` 能生成乘积网格上的所有计算结果。
 
-（这一策略可以追溯到 Matlab。）
+（这一策略可以追溯到 MATLAB。）
 
 ```{code-cell} ipython3
 grid = np.linspace(-3, 3, 3_000)
@@ -231,24 +229,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` 时，它们会相互干扰。
+每个线程都会得到自己的 `m` 私有副本，初始化为 `-np.inf`，并在其负责的迭代块中正确地更新它。
 
-线程读取了 `m` 的过时值，或者相互覆盖了更新——或者 `m` 始终保持其初始值而从未被更新。
+但在循环结束时，Numba 需要将各线程的 `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 +281,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)
@@ -320,13 +336,13 @@ with qe.Timer(precision=8):
     z_max.block_until_ready()
 ```
 
-编译完成后，JAX 明显快于 NumPy，尤其是在 GPU 上。
+编译完成后，JAX 明显快于 NumPy，在 GPU 上尤为如此。
 
 编译开销是一次性成本，当函数被反复调用时，这种开销是值得的。
 
 ### JAX 加 vmap
 
-NumPy 代码和 JAX 代码都存在一个问题：
+NumPy 代码和上述 JAX 代码都存在一个问题：
 
 虽然扁平数组占用内存较少
 
@@ -344,9 +360,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)
@@ -373,31 +389,19 @@ with qe.Timer(precision=8):
     z_max.block_until_ready()
 ```
 
-通过避免使用大型输入数组 `x_mesh` 和 `y_mesh`，这个 `vmap` 版本使用的内存少得多。
-
-在 CPU 上运行时，其运行时间与网格版本相似。
-
-在 GPU 上运行时，通常速度要快得多。
-
-实际上，使用 `vmap` 还有另一个优势：它允许我们将向量化分阶段进行。
-
-这往往会产生比传统向量化代码更易于理解的代码。
-
-当我们处理更大的问题时，将进一步探讨这些想法。
+通过避免使用大型输入数组 `x_mesh` 和 `y_mesh`，这个 `vmap` 版本使用的内存少得多，运行时间也相近。
 
-### vmap 版本 2
+但我们仍然留有一些速度提升的空间未被利用。
 
-我们可以使用 vmap 进一步提高内存效率。
+上面的代码计算了完整的二维数组 `f(x,y)`，然后再取最大值。
 
-在前一个版本中，虽然我们避免了大型输入数组，但在计算最大值之前仍然会创建大型输出数组 `f(x,y)`。
+此外，`jnp.max` 调用位于 JIT 编译函数 `f` 之外，因此编译器无法将这些操作融合为单个内核。
 
-让我们尝试一种略有不同的方法，将求最大值操作移到内部。
-
-由于这一改变，我们永远不会计算二维数组 `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))
     # 向量化该函数，以便我们可以同时对所有行调用
@@ -413,24 +417,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 v2 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 是向量化运算的赢家。
@@ -536,7 +542,7 @@ with qe.Timer(precision=8):
 
 JAX 对于这种顺序运算也相当高效。
 
-JAX 和 Numba 在编译后都能提供出色的性能，对于纯顺序运算，Numba 通常（但并非总是）提供略快的速度。
+JAX 和 Numba 在编译后都能提供出色的性能。
 
 ### 总结
 
@@ -550,30 +556,31 @@ Numba 版本简单直观，易于阅读：我们只需分配一个数组，然
 
 此外，JAX 的不可变数组意味着我们无法简单地就地更新数组元素，这使得直接复制 Numba 使用的算法变得困难。
 
-对于这类顺序运算，在代码清晰度、实现便利性以及高性能方面，Numba 是明显的赢家。
+对于这类顺序运算，在代码清晰度和实现便利性方面，Numba 是明显的赢家。
 
 ## 总体建议
 
-让我们退一步，总结一下各方案的权衡取舍。
+让我们退一步，总结一下各方的权衡取舍。
 
 对于**向量化操作**，JAX 是最强的选择。
 
-得益于 JIT 编译和跨 CPU 与 GPU 的高效并行化，它在速度上与 NumPy 持平甚至超越 NumPy。
+得益于 JIT 编译和在 CPU 与 GPU 上的高效并行化，它在速度上与 NumPy 持平或超越 NumPy。
 
-`vmap` 变换可以减少内存使用，并且通常比基于传统网格（meshgrid）的向量化方式产生更清晰的代码。
+`vmap` 变换降低了内存使用量，并且通常比传统的基于网格的向量化产生更清晰的代码。
 
-此外，JAX 函数支持自动微分，我们将在 {doc}`autodiff` 中进行探讨。
+此外，JAX 函数支持自动微分，我们将在 {doc}`autodiff` 中进一步探讨。
 
 对于**顺序操作**，Numba 具有明显优势。
 
-代码自然易读——只需一个带装饰器的 Python 循环——且性能出色。
+代码自然且可读——只需一个带有装饰器的 Python 循环——性能也非常出色。
 
-JAX 可以通过 `lax.scan` 处理顺序问题，但对于纯顺序工作而言，其语法不够直观，性能提升也十分有限。
+JAX 可以通过 `lax.scan` 处理顺序问题，但语法不够直观。
 
-话虽如此，`lax.scan` 有一个重要优势：它支持对循环进行自动微分，而 Numba 无法做到这一点。
-
-如果需要对顺序计算进行微分（例如，计算轨迹对模型参数的敏感性），尽管语法不够自然，JAX 仍是更好的选择。
+```{note}
+`lax.scan` 的一个重要优势是它支持通过循环进行自动微分，而 Numba 无法做到这一点。
+如果您需要对顺序计算进行微分（例如，计算轨迹对模型参数的敏感性），尽管语法不够自然，JAX 仍是更好的选择。
+```
 
-在实践中，许多问题往往同时涉及两种模式。
+在实践中，许多问题涉及两种模式的混合。
 
-一个实用的经验法则是：新项目默认使用 JAX，尤其是在硬件加速或可微分性可能有用的情况下；当需要一个快速且可读的紧凑顺序循环时，则选用 Numba。
+一个实用的经验法则是：对于新项目默认使用 JAX，尤其是当硬件加速或可微分性可能有用时，而当您有一个需要快速且可读的紧凑顺序循环时，则选用 Numba。