Update translation: lectures/numpy_vs_numba_vs_jax.md

Author: mmcky

lectures/numpy_vs_numba_vs_jax.md

@@ -24,6 +24,7 @@ translation:
     Sequential operations::Numba Version: Numba 版本
     Sequential operations::JAX Version: JAX 版本
     Sequential operations::Summary: 总结
+    Overall recommendations: 总体建议
 ---
 
 (parallel)=
@@ -69,7 +70,10 @@ tags: [hide-output]
 
 ```{code-cell} ipython3
 import random
+from functools import partial
+
 import numpy as np
+import numba
 import quantecon as qe
 import matplotlib.pyplot as plt
 import matplotlib as mpl  # i18n
@@ -80,6 +84,7 @@ from mpl_toolkits.mplot3d.axes3d import Axes3D
 from matplotlib import cm
 import jax
 import jax.numpy as jnp
+from jax import lax
 ```
 
 ## 向量化运算
@@ -117,7 +122,7 @@ ax.plot_surface(x,
                 y,
                 f(x, y),
                 rstride=2, cstride=2,
-                cmap=cm.jet,
+                cmap=cm.viridis,
                 alpha=0.7,
                 linewidth=0.25)
 ax.set_zlim(-0.5, 1.0)
@@ -143,7 +148,6 @@ for x in grid:
             m = z
 ```
 
-
 ### NumPy 向量化
 
 如果我们切换到 NumPy 风格的向量化，就可以使用更大的网格，并且代码执行速度相对较快。
@@ -168,14 +172,11 @@ print(f"NumPy result: {z_max_numpy:.6f}")
 
 （并行化效率不高，因为二进制文件在看到数组 `x` 和 `y` 的大小之前就已经被编译了。）
 
-
 ### 与 Numba 的比较
 
 现在让我们看看能否使用简单循环的 Numba 获得更好的性能。
 
 ```{code-cell} ipython3
-import numba
-
 @numba.jit
 def compute_max_numba(grid):
     m = -np.inf
@@ -189,9 +190,9 @@ def compute_max_numba(grid):
 grid = np.linspace(-3, 3, 3_000)
 
 with qe.Timer(precision=8):
-    z_max_numpy = compute_max_numba(grid)
+    z_max_numba = compute_max_numba(grid)
 
-print(f"Numba result: {z_max_numpy:.6f}")
+print(f"Numba result: {z_max_numba:.6f}")
 ```
 
 让我们再次运行以消除编译时间。
@@ -207,7 +208,6 @@ with qe.Timer(precision=8):
 
 另一方面，Numba 例程使用的内存少得多，因为我们只处理一个一维网格。
 
-
 ### 并行化的 Numba
 
 现在让我们使用 `prange` 尝试 Numba 的并行化：
@@ -282,7 +282,6 @@ with qe.Timer(precision=8):
 
 对于更强大的机器和更大的网格尺寸，即使在 CPU 上，并行化也能带来显著的速度提升。
 
-
 ### 使用 JAX 的向量化代码
 
 表面上，JAX 中的向量化代码与 NumPy 代码类似。
@@ -303,7 +302,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))
@@ -320,11 +319,10 @@ with qe.Timer(precision=8):
     z_max.block_until_ready()
 ```
 
-编译完成后，由于 GPU 加速，JAX 明显快于 NumPy。
+编译完成后，JAX 明显快于 NumPy，尤其是在 GPU 上。
 
 编译开销是一次性成本，当函数被反复调用时，这种开销是值得的。
 
-
 ### JAX 加 vmap
 
 NumPy 代码和 JAX 代码都存在一个问题：
@@ -386,7 +384,6 @@ with qe.Timer(precision=8):
 
 当我们处理更大的问题时，将进一步探讨这些想法。
 
-
 ### vmap 版本 2
 
 我们可以使用 vmap 进一步提高内存效率。
@@ -421,7 +418,7 @@ def compute_max_vmap_v2(grid):
 with qe.Timer(precision=8):
     z_max = compute_max_vmap_v2(grid).block_until_ready()
 
-print(f"JAX vmap v1 result: {z_max:.6f}")
+print(f"JAX vmap v2 result: {z_max:.6f}")
 ```
 
 让我们再次运行以消除编译时间：
@@ -433,7 +430,6 @@ with qe.Timer(precision=8):
 
 如果您像我们一样在 GPU 上运行，应该能看到又一个不小的速度提升。
 
-
 ### 总结
 
 在我们看来，JAX 是向量化运算的赢家。
@@ -448,15 +444,13 @@ with qe.Timer(precision=8):
 
 对于经济学、计量经济学和金融学中遇到的大多数情况，将高效并行化的工作交给 JAX 编译器，远比尝试手工编写这些例程要好得多。
 
-
 ## 顺序运算
 
 某些运算本质上是顺序的——因此难以或不可能向量化。
 
 在这种情况下，NumPy 是一个较差的选择，我们只剩下 Numba 或 JAX 可以选择。
 
-为了比较这两种选择，我们将重新回顾在{doc}`Numba 讲座 <numba>`中看到的迭代二次映射问题。
-
+为了比较这两种选择，我们将重新回顾在 {doc}`Numba 讲座 <numba>` 中看到的迭代二次映射问题。
 
 ### Numba 版本
 
@@ -501,9 +495,6 @@ Numba 的编译通常相当快，对于像这样的顺序运算，生成的代
 （我们将 `n` 设为静态，因为它影响数组大小，JAX 希望在编译代码中针对其值进行特化处理。）
 
 ```{code-cell} ipython3
-from jax import lax
-from functools import partial
-
 cpu = jax.devices("cpu")[0]
 
 @partial(jax.jit, static_argnums=(1,), device=cpu)
@@ -546,7 +537,6 @@ JAX 对于这种顺序运算也相当高效。
 
 JAX 和 Numba 在编译后都能提供出色的性能，对于纯顺序运算，Numba 通常（但并非总是）提供略快的速度。
 
-
 ### 总结
 
 虽然 Numba 和 JAX 在顺序运算中都能提供出色的性能，但**在代码可读性和易用性方面存在显著差异**。
@@ -559,4 +549,30 @@ Numba 版本简单直观，易于阅读：我们只需分配一个数组，然
 
 此外，JAX 的不可变数组意味着我们无法简单地就地更新数组元素，这使得直接复制 Numba 使用的算法变得困难。
 
-对于这类顺序运算，在代码清晰度、实现便利性以及高性能方面，Numba 是明显的赢家。
+对于这类顺序运算，在代码清晰度、实现便利性以及高性能方面，Numba 是明显的赢家。
+
+## 总体建议
+
+让我们退一步，总结一下各方案的权衡取舍。
+
+对于**向量化操作**，JAX 是最强的选择。
+
+得益于 JIT 编译和跨 CPU 与 GPU 的高效并行化，它在速度上与 NumPy 持平甚至超越 NumPy。
+
+`vmap` 变换可以减少内存使用，并且通常比基于传统网格（meshgrid）的向量化方式产生更清晰的代码。
+
+此外，JAX 函数支持自动微分，我们将在 {doc}`autodiff` 中进行探讨。
+
+对于**顺序操作**，Numba 具有明显优势。
+
+代码自然易读——只需一个带装饰器的 Python 循环——且性能出色。
+
+JAX 可以通过 `lax.scan` 处理顺序问题，但对于纯顺序工作而言，其语法不够直观，性能提升也十分有限。
+
+话虽如此，`lax.scan` 有一个重要优势：它支持对循环进行自动微分，而 Numba 无法做到这一点。
+
+如果需要对顺序计算进行微分（例如，计算轨迹对模型参数的敏感性），尽管语法不够自然，JAX 仍是更好的选择。
+
+在实践中，许多问题往往同时涉及两种模式。
+
+一个实用的经验法则是：新项目默认使用 JAX，尤其是在硬件加速或可微分性可能有用的情况下；当需要一个快速且可读的紧凑顺序循环时，则选用 Numba。