Update translation: lectures/numpy_vs_numba_vs_jax.md

Author: mmcky

lectures/numpy_vs_numba_vs_jax.md

@@ -17,17 +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: 总结
+    Vectorized operations::Summary: vmap 版本 2
     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">
@@ -69,7 +67,6 @@ tags: [hide-output]
 我们将使用以下导入。
 
 ```{code-cell} ipython3
-import random
 from functools import partial
 
 import numpy as np
@@ -472,14 +469,16 @@ def qm(x0, n, α=4.0):
 ```{code-cell} ipython3
 n = 10_000_000
 
-with qe.Timer(precision=8):
+with qe.Timer():
+    # First run
     x = qm(0.1, n)
 ```
 
 让我们再次运行以消除编译时间：
 
 ```{code-cell} ipython3
-with qe.Timer(precision=8):
+with qe.Timer():
+    # Second run
     x = qm(0.1, n)
 ```
 
@@ -491,15 +490,62 @@ Numba 的编译通常相当快，对于像这样的顺序运算，生成的代
 
 ### JAX 版本
 
-现在让我们使用 `lax.scan` 创建一个 JAX 版本：
+现在让我们使用 `at[t].set` 风格的语法创建一个 JAX 版本，正如 {ref}`JAX 讲座中讨论的 <jax_at_workaround>`，这为不可变数组提供了一种变通方法。
 
-（我们将 `n` 设为静态，因为它影响数组大小，JAX 希望在编译代码中针对其值进行特化处理。）
+我们将使用 `lax.fori_loop`，它是一种可以被 XLA 编译的 for 循环版本。
 
 ```{code-cell} ipython3
 cpu = jax.devices("cpu")[0]
 
-@partial(jax.jit, static_argnums=(1,), device=cpu)
-def qm_jax(x0, n, α=4.0):
+@partial(jax.jit, static_argnames=("n",), device=cpu)
+def qm_jax_fori(x0, n, α=4.0):
+
+    x = jnp.empty(n + 1).at[0].set(x0)
+
+    def update(t, x):
+        return x.at[t + 1].set(α * x[t] * (1 - x[t]))
+
+    x = lax.fori_loop(0, n, update, x)
+    return x
+
+```
+
+* 我们将 `n` 设为静态，因为它影响数组大小，JAX 希望在编译代码中针对其值进行特化处理。
+* 我们通过 `device=cpu` 将计算固定在 CPU 上，因为这种顺序工作负载由许多小操作组成，几乎没有机会利用 GPU 并行性。
+
+虽然 `at[t].set` 看起来在每一步都创建了一个新数组，但在 JIT 编译的函数内部，编译器会检测到旧数组不再需要，并就地执行更新。
+
+让我们使用相同的参数计时：
+
+```{code-cell} ipython3
+with qe.Timer():
+    # First run
+    x_jax = qm_jax_fori(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
+```
+
+让我们再次运行以消除编译开销：
+
+```{code-cell} ipython3
+with qe.Timer():
+    # Second run
+    x_jax = qm_jax_fori(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
+```
+
+JAX 对于这种顺序运算也相当高效。
+
+
+我们还有另一种实现循环的方式，使用 `lax.scan`。
+
+这种替代方案可以说更符合 JAX 的函数式方法——尽管语法难以记忆。
+
+
+```{code-cell} ipython3
+@partial(jax.jit, static_argnames=("n",), device=cpu)
+def qm_jax_scan(x0, n, α=4.0):
     def update(x, t):
         x_new = α * x * (1 - x)
         return x_new, x_new
@@ -510,33 +556,27 @@ def qm_jax(x0, n, α=4.0):
 
 这段代码不易阅读，但本质上，`lax.scan` 反复调用 `update` 并将返回值 `x_new` 累积到一个数组中。
 
-```{note}
-细心的读者会注意到，我们在 `jax.jit` 装饰器中指定了 `device=cpu`。
-
-该计算由许多小的顺序运算组成，几乎没有机会让 GPU 利用并行性。
-
-因此，GPU 上的内核启动开销往往占主导地位，使得 CPU 更适合这种工作负载。
-
-好奇的读者可以尝试删除此选项，看看性能如何变化。
-```
-
 让我们使用相同的参数计时：
 
 ```{code-cell} ipython3
-with qe.Timer(precision=8):
-    x_jax = qm_jax(0.1, n).block_until_ready()
+with qe.Timer():
+    # First run
+    x_jax = qm_jax_scan(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
 ```
 
 让我们再次运行以消除编译开销：
 
 ```{code-cell} ipython3
-with qe.Timer(precision=8):
-    x_jax = qm_jax(0.1, n).block_until_ready()
+with qe.Timer():
+    # Second run
+    x_jax = qm_jax_scan(0.1, n)
+    # Hold interpreter
+    x_jax.block_until_ready()
 ```
 
-JAX 对于这种顺序运算也相当高效。
-
-JAX 和 Numba 在编译后都能提供出色的性能，对于纯顺序运算，Numba 通常（但并非总是）提供略快的速度。
+JAX 和 Numba 在编译后都能提供出色的性能。
 
 ### 总结
 
@@ -546,11 +586,11 @@ Numba 版本简单直观，易于阅读：我们只需分配一个数组，然
 
 这正是大多数程序员思考该算法的方式。
 
-另一方面，JAX 版本需要使用 `lax.scan`，这明显不够直观。
+另一方面，JAX 版本需要使用 `lax.fori_loop` 或 `lax.scan`，两者都比标准 Python 循环更不直观。
 
-此外，JAX 的不可变数组意味着我们无法简单地就地更新数组元素，这使得直接复制 Numba 使用的算法变得困难。
+虽然 JAX 的 `at[t].set` 语法确实允许逐元素更新，但整体代码仍然比 Numba 等价版本更难阅读。
 
-对于这类顺序运算，在代码清晰度、实现便利性以及高性能方面，Numba 是明显的赢家。
+对于这类顺序运算，在代码清晰度和实现便利性方面，Numba 是明显的赢家。
 
 ## 总体建议
 
@@ -568,11 +608,12 @@ Numba 版本简单直观，易于阅读：我们只需分配一个数组，然
 
 代码自然易读——只需一个带装饰器的 Python 循环——且性能出色。
 
-JAX 可以通过 `lax.scan` 处理顺序问题，但对于纯顺序工作而言，其语法不够直观，性能提升也十分有限。
-
-话虽如此，`lax.scan` 有一个重要优势：它支持对循环进行自动微分，而 Numba 无法做到这一点。
+JAX 可以通过 `lax.fori_loop` 或 `lax.scan` 处理顺序问题，但语法不够直观。
 
+```{note}
+`lax.fori_loop` 和 `lax.scan` 有一个重要优势：它们支持对循环进行自动微分，而 Numba 无法做到这一点。
 如果需要对顺序计算进行微分（例如，计算轨迹对模型参数的敏感性），尽管语法不够自然，JAX 仍是更好的选择。
+```
 
 在实践中，许多问题往往同时涉及两种模式。