diff --git a/.translate/state/mix_model.md.yml b/.translate/state/mix_model.md.yml
new file mode 100644
index 00000000..a8511c7b
--- /dev/null
+++ b/.translate/state/mix_model.md.yml
@@ -0,0 +1,6 @@
+source-sha: ebb81e76e7138ac63f7cfde4142dd5a062d8004b
+synced-at: "2026-07-21"
+model: claude-sonnet-5
+mode: RESYNC
+section-count: 8
+tool-version: 0.20.0
diff --git a/lectures/mix_model.md b/lectures/mix_model.md
index 4ddff447..a250cb35 100644
--- a/lectures/mix_model.md
+++ b/lectures/mix_model.md
@@ -9,6 +9,17 @@ kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
   name: python3
+translation:
+  title: 错误模型
+  headings:
+    Overview: 概述
+    Sampling from compound lottery $H$: 从复合彩票 $H$ 中抽样
+    Type 1 agent: 类型1智能体
+    What a type 1 agent learns when mixture $H$ generates data: 当混合分布$H$生成数据时，类型1代理学到什么
+    Kullback-Leibler divergence governs limit of $\pi_t$: Kullback-Leibler 散度决定 $\pi_t$ 的极限
+    Type 2 agent: 类型2代理
+    Concluding remarks: 总结性评论
+    Exercises: 练习
 ---
 
 (likelihood-ratio-process)=
@@ -20,12 +31,12 @@ kernelspec:
 ```{code-cell} ipython3
 :tags: [no-execute, hide-output]
 
-!pip install numpyro jax
+!pip install numpyro
 ```
 
 ## 概述
 
-这是{doc}`这个 quantecon 讲座 <likelihood_bayes>`的续篇。
+这是[](likelihood_bayes)的续篇。
 
 我们将讨论创建复合彩票的两种方式及其后果。
 
@@ -37,7 +48,7 @@ kernelspec:
 
 * 另一种方式是在每个时期都在相同的两个可能的概率分布之间进行混合
 
-这个统计设定与那个 quantecon 讲座中研究的问题相近但不完全相同。
+这个统计设定与那个讲座中研究的问题相近但不完全相同。
 
 在那个讲座中,有两个可能支配非负随机变量 $W$ 连续抽取的独立同分布过程。
 
@@ -52,15 +63,14 @@ kernelspec:
 我们研究了代理人如何使用条件概率法则和观察到的历史 $w^t =\{w_s\}_{s=0}^t$ 来形成
 
 $$
-
-\pi_t = E [ \textrm{自然选择分布} f | w^t] , \quad  t = 0, 1, 2, \ldots
+\pi_t = \mathbb{E}[ \textrm{自然选择分布} f \mid w^t] , \quad  t = 0, 1, 2, \ldots
 $$
 
 然而，在本讲座的设定中，该规则为智能体赋予了一个错误的模型。
 
 原因是现在工资序列实际上是由一个不同的统计模型描述的。
 
-因此，我们需要以下列方式改变{doc}`quantecon讲座<likelihood_bayes>`的规范。
+因此，我们需要以下列方式改变[](likelihood_bayes)中的设定。
 
 现在，在**每个时期**$t \geq 0$，自然投掷一个可能不公平的硬币，以概率$\alpha$出现$f$，以概率$1-\alpha$出现$g$。
 
@@ -78,7 +88,7 @@ $$
 
 我们的第二类智能体正确地知道，自然在每个时期以混合概率$\alpha \in (0,1)$混合$f$和$g$，尽管智能体不知道混合参数。
 
-我们的第一类智能体应用{doc}`这个quantecon讲座<likelihood_bayes>`中描述的学习算法。
+我们的第一类智能体应用[](likelihood_bayes)中描述的学习算法。
 
 在该讲座中所涉及的统计模型的背景下，这是一个好的学习算法，它使贝叶斯学习者能够
 
@@ -106,7 +116,7 @@ $$
 
 * 自然如何在两个分布 $f$ 和 $g$ 之间进行*混合*以创建一个新的分布 $h$。
 
-* Kullback-Leibler统计散度 <https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence> 用于控制在不正确统计模型下的统计学习
+* [Kullback-Leibler统计散度](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence)，用于控制在不正确统计模型下的统计学习
 
 * 一个有用的Python函数`numpy.searchsorted`，它与均匀随机数生成器结合使用时，可以用来从任意分布中进行采样
 
@@ -116,38 +126,26 @@ $$
 :hide-output: false
 
 import matplotlib.pyplot as plt
-import matplotlib as mpl
-FONTPATH = "fonts/SourceHanSerifSC-SemiBold.otf"
-import matplotlib as mpl
-mpl.font_manager.fontManager.addfont(FONTPATH)
-plt.rcParams['font.family'] = ['Source Han Serif SC']
-
+import seaborn as sns
 import numpy as np
-from numba import vectorize, jit
-from math import gamma
-import pandas as pd
-import scipy.stats as sp
 from scipy.integrate import quad
 
-import seaborn as sns
-colors = sns.color_palette()
-
 import numpyro
+import jax
 import numpyro.distributions as dist
 from numpyro.infer import MCMC, NUTS
-
 import jax.numpy as jnp
+import jax.scipy.stats as jsp
 from jax import random
+from jax.scipy.special import gamma
 
-np.random.seed(0)
+# 使JAX能够使用64位浮点运算
+jax.config.update("jax_enable_x64", True)
 
-@jit
-def set_seed():
-    np.random.seed(0)
-set_seed()
+colors = sns.color_palette()
 ```
 
-让我们用Python生成两个贝塔分布
+让我们用Python生成两个贝塔分布。
 
 ```{code-cell} ipython3
 :hide-output: false
@@ -156,51 +154,46 @@ set_seed()
 F_a, F_b = 1, 1
 G_a, G_b = 3, 1.2
 
-@vectorize
+
+@jax.jit
 def p(x, a, b):
     r = gamma(a + b) / (gamma(a) * gamma(b))
-    return r * x** (a-1) * (1 - x) ** (b-1)
+    return r * x**(a-1) * (1 - x)**(b-1)
 
 # 两个密度函数
-f = jit(lambda x: p(x, F_a, F_b))
-g = jit(lambda x: p(x, G_a, G_b))
+f = jax.jit(lambda x: p(x, F_a, F_b))
+g = jax.jit(lambda x: p(x, G_a, G_b))
 ```
 
 ```{code-cell} ipython3
 :hide-output: false
 
-@jit
-def simulate(a, b, T=50, N=500):
-    '''
+def simulate(key, a, b, T=50, N=500):
+    """
     生成N组T个似然比观测值，
     以N x T矩阵形式返回。
-
-    '''
-
-    l_arr = np.empty((N, T))
-
-    for i in range(N):
-
-        for j in range(T):
-            w = np.random.beta(a, b)
-            l_arr[i, j] = f(w) / g(w)
-
+    """
+    w = jax.random.beta(key, a, b, (N, T))
+    l_arr = f(w) / g(w)
     return l_arr
 ```
 
-我们还将使用以下Python代码来准备一些信息丰富的模拟
+我们还将使用以下Python代码来准备一些信息丰富的模拟。
 
 ```{code-cell} ipython3
 :hide-output: false
 
-l_arr_g = simulate(G_a, G_b, N=50000)
+key = jax.random.key(42)
+key_g, key_f = jax.random.split(key)
+
+l_arr_g = simulate(key_g, G_a, G_b, N=50000)
 l_seq_g = np.cumprod(l_arr_g, axis=1)
 ```
 
 ```{code-cell} ipython3
 :hide-output: false
 
-l_arr_f = simulate(F_a, F_b, N=50000)
+l_arr_f = simulate(key_f, F_a, F_b, N=50000)
 l_seq_f = np.cumprod(l_arr_f, axis=1)
 ```
 
@@ -214,7 +207,7 @@ l_seq_f = np.cumprod(l_arr_f, axis=1)
 
 * 第一步：
 
-  * 使用 numpy.random.choice 函数抛一个不公平的硬币，以概率 $\alpha$ 选择分布 $F$，以概率 $1-\alpha$ 选择 $G$
+  * 使用均匀随机抽取来抛一个不公平的硬币，以概率 $\alpha$ 选择分布 $F$，以概率 $1-\alpha$ 选择 $G$
 
 * 第二步：
 
@@ -225,7 +218,7 @@ l_seq_f = np.cumprod(l_arr_f, axis=1)
   * 将前两步放在一个大循环中，对 $w$ 的每个实现都执行这些步骤
 
 
-我们的第二种方法使用均匀分布和以下在 quantecon 讲座 <https://python.quantecon.org/prob_matrix.html> 中描述和使用的事实：
+我们的第二种方法使用均匀分布和以下我们在[](prob_matrix)中也描述和使用过的事实：
 
  * 如果随机变量 $X$ 的累积分布函数为 $F$，那么随机变量 $F^{-1}(U)$ 也具有累积分布函数 $F$，其中 $U$ 是 $[0,1]$ 上的均匀随机变量。
 
@@ -233,50 +226,63 @@ l_seq_f = np.cumprod(l_arr_f, axis=1)
 
 我们将结合 `numpy.searchsorted` 命令使用这个事实来直接从 $H$ 中抽样。
 
-关于 `searchsorted` 函数的说明，请参见 <https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html>。
+关于 `searchsorted` 函数的详细说明，请参见[`numpy.searchsorted` 文档](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html)。
 
 观看[Mr. P Solver关于蒙特卡洛模拟的视频](https://www.google.com/search?q=Mr.+P+Solver+video+on+Monte+Carlo+simulation&oq=Mr.+P+Solver+video+on+Monte+Carlo+simulation)，了解这个强大技巧的其他应用。
 
 在下面的Python代码中，我们将使用两种方法，并确认它们都能很好地从我们的目标混合分布中进行采样。
 
 ```{code-cell} ipython3
-@jit
-def draw_lottery(p, N):
-    "直接从复合彩票中抽取。"
-
-    draws = []
-    for i in range(0, N):
-        if np.random.rand()<=p:
-            draws.append(np.random.beta(F_a, F_b))
-        else:
-            draws.append(np.random.beta(G_a, G_b))
-    return np.array(draws)
-
-def draw_lottery_MC(p, N):
-    "使用蒙特卡洛技巧从复合彩票中抽取。"
-
-    xs = np.linspace(1e-8,1-(1e-8),10000)
-    CDF = p*sp.beta.cdf(xs, F_a, F_b) + (1-p)*sp.beta.cdf(xs, G_a, G_b)
-
-    Us = np.random.rand(N)
-    draws = xs[np.searchsorted(CDF[:-1], Us)]
+def draw_lottery(key, p, N):
+    """Draw from the compound lottery directly."""
+    # split the key: one for the coin flip, one for each distribution
+    key_uniform, key_beta1, key_beta2 = jax.random.split(key, 3)
+
+    u = jax.random.uniform(key_uniform, shape=(N,))
+    f_draws = jax.random.beta(key_beta1, F_a, F_b, shape=(N,))
+    g_draws = jax.random.beta(key_beta2, G_a, G_b, shape=(N,))
+
+    # select F with probability p, else G
+    return jnp.where(u <= p, f_draws, g_draws)
+
+
+# JIT-compile with `N` as a static argument
+draw_lottery = jax.jit(draw_lottery, static_argnums=(2,))
+```
+
+```{code-cell} ipython3
+def draw_lottery_MC(key, p, N):
+    "Draw from the compound lottery using the Monte Carlo trick."
+
+    xs = jnp.linspace(1e-8, 1 - (1e-8), 10000)
+    CDF = p * jsp.beta.cdf(xs, F_a, F_b) + (1 - p) * jsp.beta.cdf(xs, G_a, G_b)
+
+    Us = jax.random.uniform(key, shape=(N,))
+    draws = xs[jnp.searchsorted(CDF[:-1], Us)]
     return draws
 ```
 
 ```{code-cell} ipython3
-# 验证
+---
+mystnb:
+  figure:
+    caption: 直接抽样与蒙特卡洛抽样
+    name: fig-lottery-draws
+---
+# verify
 N = 100000
 α = 0.0
 
-sample1 = draw_lottery(α, N)
-sample2 = draw_lottery_MC(α, N)
+key1, key2 = jax.random.split(jax.random.key(42))
+sample1 = draw_lottery(key1, α, N)
+sample2 = draw_lottery_MC(key2, α, N)
 
-# 绘制抽样和密度函数
-plt.hist(sample1, 50, density=True, alpha=0.5, label='直接抽样')
-plt.hist(sample2, 50, density=True, alpha=0.5, label='MC抽样')
+# plot draws and density function
+plt.hist(sample1, 50, density=True, alpha=0.5, label="direct draws")
+plt.hist(sample2, 50, density=True, alpha=0.5, label="MC draws")
 
-xs = np.linspace(0,1,1000)
-plt.plot(xs, α*f(xs)+(1-α)*g(xs), color='red', label='密度')
+xs = np.linspace(0, 1, 1000)
+plt.plot(xs, α * f(xs) + (1 - α) * g(xs), color="red", label="density", lw=2)
 
 plt.legend()
 plt.show()
@@ -284,18 +290,18 @@ plt.show()
 
 ## 类型1智能体
 
-现在我们来研究类型1智能体学到了什么
+现在我们来研究类型1智能体学到了什么。
 
 请记住，我们的类型1智能体使用了错误的统计模型，认为自然在时间-1时就一次性地在$f$和$g$之间做出了选择。
 
-因此，类型1智能体使用了在{doc}`这个quantecon讲座<likelihood_bayes>`中研究的学习算法。
+因此，类型1智能体使用了在[](likelihood_bayes)中研究的学习算法。
 
 我们现在简要回顾一下该学习算法。
 
 让$\pi_t$表示贝叶斯后验概率，定义为
 
 $$
-\pi_t = {\rm Prob}(q=f|w^t)
+\pi_t = \mathbb{P}\{q = f \mid w^t\}
 $$
 
 似然比过程在控制后验概率$\pi_t$演化的公式中起着主要作用，这是**贝叶斯法则**的一个实例。
@@ -313,13 +319,11 @@ $$ (eq:recur1)
 ```{code-cell} ipython3
 :hide-output: false
 
-@jit
+@jax.jit
 def update(π, l):
-    "Update π using likelihood l"
-
+    """Update π using likelihood l"""
     # Update belief
     π = π * l / (π * l + 1 - π)
-
     return π
 ```
 
@@ -387,64 +391,87 @@ $$ (eq:bayeslaw103)
 让我们设定一个$\alpha$值，然后观察$\pi_t$如何演变。
 
 ```{code-cell} ipython3
-def simulate_mixed(α, T=50, N=500):
+def simulate_mixed(key, α, T=50, N=500):
     """
-    当真实密度为混合h;α时，生成N组T个似然比观测值，
-    返回N x T矩阵
+    Generate N sets of T observations of the likelihood ratio,
+    return as N x T matrix, when the true density is mixed h;α
     """
 
-    w_s = draw_lottery(α, N*T).reshape(N, T)
+    w_s = draw_lottery(key, α, N * T).reshape(N, T)
     l_arr = f(w_s) / g(w_s)
-
     return l_arr
 
-def plot_π_seq(α, π1=0.2, π2=0.8, T=200):
+
+def plot_π_seq(key, α, π1=0.2, π2=0.8, T=200):
     """
-    当混合分布支配数据时，计算并绘制π_seq和对数似然比过程
+    Compute and plot π_seq and the log likelihood ratio process
+    when the mixed distribution governs the data.
     """
 
-    l_arr_mixed = simulate_mixed(α, T=T, N=50)
+    l_arr_mixed = simulate_mixed(key, α, T=T, N=50)
     l_seq_mixed = np.cumprod(l_arr_mixed, axis=1)
 
     T = l_arr_mixed.shape[1]
-    π_seq_mixed = np.empty((2, T+1))
+    π_seq_mixed = np.empty((2, T + 1))
     π_seq_mixed[:, 0] = π1, π2
 
     for t in range(T):
         for i in range(2):
-            π_seq_mixed[i, t+1] = update(π_seq_mixed[i, t], l_arr_mixed[0, t])
+            π_seq_mixed[i, t + 1] = update(
+                π_seq_mixed[i, t], l_arr_mixed[0, t]
+            )
 
-    # 绘图
+    # plot
     fig, ax1 = plt.subplots()
     for i in range(2):
-        ax1.plot(range(T+1), π_seq_mixed[i, :], label=rf"$\pi_0$={π_seq_mixed[i, 0]}")
-
-    ax1.plot(np.nan, np.nan,  '--', color='b', label='对数似然比过程')
+        ax1.plot(
+            range(T + 1),
+            π_seq_mixed[i, :],
+            label=rf"$\pi_0$={π_seq_mixed[i, 0]}",
+            lw=2,
+        )
+
+    ax1.plot(
+        np.nan, np.nan, "--", color="b",
+        label="Log likelihood ratio process", lw=2,
+    )
     ax1.set_ylabel(r"$\pi_t$")
-    ax1.set_xlabel("t")
+    ax1.set_xlabel("$t$")
     ax1.legend()
-    ax1.set_title("当$\\alpha F + (1-\\alpha)G$支配数据时")
 
     ax2 = ax1.twinx()
-    ax2.plot(range(1, T+1), np.log(l_seq_mixed[0, :]), '--', color='b')
-    ax2.set_ylabel("$log(L(w^{t}))$")
+    ax2.plot(range(1, T + 1), np.log(l_seq_mixed[0, :]), "--", color="b", lw=2)
+    ax2.set_ylabel(r"$\log(L(w^{t}))$")
 
     plt.show()
 ```
 
 ```{code-cell} ipython3
-plot_π_seq(α = 0.6)
+---
+mystnb:
+  figure:
+    caption: 信念路径，$\alpha = 0.6$
+    name: fig-pi-seq-1
+---
+plot_π_seq(jax.random.key(42), α=0.6)
 ```
 
 上图显示了对数似然比过程的一个样本路径（蓝色虚线），以及从两个不同初始条件开始的 $\pi_t$ 的样本路径。
 
-让我们看看当我们改变 $\alpha$ 时会发生什么
+
+让我们看看当我们改变 $\alpha$ 时会发生什么。
 
 ```{code-cell} ipython3
-plot_π_seq(α = 0.2)
+---
+mystnb:
+  figure:
+    caption: 信念路径，$\alpha = 0.2$
+    name: fig-pi-seq-2
+---
+plot_π_seq(jax.random.key(42), α=0.2)
 ```
 
-显然，$\alpha$ 对 $\pi_t$ 在 $t \rightarrow + \infty$ 时的终点有很大影响
+显然，$\alpha$ 对 $\pi_t$ 在 $t \rightarrow + \infty$ 时的终点有很大影响。
 
 ## Kullback-Leibler 散度决定 $\pi_t$ 的极限
 
@@ -475,78 +502,98 @@ $$ \min_{f,g} \{KL_g, KL_f\} $$
 当且仅当 $KL_f < KL_g$ 时，$\pi_t$ 在 $t \rightarrow +\infty$ 时趋向于 1。
 
 ```{code-cell} ipython3
-@vectorize
+@jax.jit
 def KL_g(α):
     "Compute the KL divergence KL(h, g)."
-    err = 1e-8                          # to avoid 0 at end points
-    ws = np.linspace(err, 1-err, 10000)
+    err = 1e-8  # to avoid 0 at end points
+    ws = jnp.linspace(err, 1 - err, 10000)
     gs, fs = g(ws), f(ws)
-    hs = α*fs + (1-α)*gs
-    return np.sum(np.log(hs/gs)*hs)/10000
+    hs = α * fs + (1 - α) * gs
+    return jnp.sum(jnp.log(hs / gs) * hs) / 10000
+
+
+KL_g_v = jax.vmap(KL_g)
 
-@vectorize
+
+@jax.jit
 def KL_f(α):
     "Compute the KL divergence KL(h, f)."
-    err = 1e-8                          # to avoid 0 at end points
-    ws = np.linspace(err, 1-err, 10000)
+    err = 1e-8  # to avoid 0 at end points
+    ws = jnp.linspace(err, 1 - err, 10000)
     gs, fs = g(ws), f(ws)
-    hs = α*fs + (1-α)*gs
-    return np.sum(np.log(hs/fs)*hs)/10000
+    hs = α * fs + (1 - α) * gs
+    return jnp.sum(jnp.log(hs / fs) * hs) / 10000
+
+
+KL_f_v = jax.vmap(KL_f)
 
 
 # compute KL using quad in Scipy
 def KL_g_quad(α):
     "Compute the KL divergence KL(h, g) using scipy.integrate."
-    h = lambda x: α*f(x) + (1-α)*g(x)
-    return quad(lambda x: h(x) * np.log(h(x)/g(x)), 0, 1)[0]
+    h = lambda x: α * f(x) + (1 - α) * g(x)
+    return quad(lambda x: h(x) * np.log(h(x) / g(x)), 0, 1)[0]
+
 
 def KL_f_quad(α):
     "Compute the KL divergence KL(h, f) using scipy.integrate."
-    h = lambda x: α*f(x) + (1-α)*g(x)
-    return quad(lambda x: h(x) * np.log(h(x)/f(x)), 0, 1)[0]
+    h = lambda x: α * f(x) + (1 - α) * g(x)
+    return quad(lambda x: h(x) * np.log(h(x) / f(x)), 0, 1)[0]
+
 
 # vectorize
 KL_g_quad_v = np.vectorize(KL_g_quad)
 KL_f_quad_v = np.vectorize(KL_f_quad)
 
 
-# Let us find the limit point
-def π_lim(α, T=5000, π_0=0.4):
-    "Find limit of π sequence."
-    π_seq = np.zeros(T+1)
-    π_seq[0] = π_0
-    l_arr = simulate_mixed(α, T, N=1)[0]
+@jax.jit
+def π_lim(key, α, T=5000, π_0=0.4):
+    """Find limit of π sequence."""
+    # Get lottery draws
+    l_arr = simulate_mixed(key, α, T, N=1)[0]
 
-    for t in range(T):
-        π_seq[t+1] = update(π_seq[t], l_arr[t])
-    return π_seq[-1]
+    def scan_fn(π_prev, l_t):
+        π_new = update(π_prev, l_t)
+        return π_new, π_new
+
+    # Scan over lottery draws
+    π_final, _ = jax.lax.scan(scan_fn, π_0, l_arr)
 
-π_lim_v = np.vectorize(π_lim)
+    return π_final
+
+
+π_lim_v = jax.vmap(π_lim)
 ```
 
 让我们首先绘制每个$\alpha$对应的KL散度$KL_g\left(\alpha\right), KL_f\left(\alpha\right)$。
 
 ```{code-cell} ipython3
+---
+mystnb:
+  figure:
+    caption: KL散度与$\alpha$的关系
+    name: fig-kl
+---
 α_arr = np.linspace(0, 1, 100)
-KL_g_arr = KL_g(α_arr)
-KL_f_arr = KL_f(α_arr)
+KL_g_arr = KL_g_v(α_arr)
+KL_f_arr = KL_f_v(α_arr)
 
-fig, ax = plt.subplots(1, figsize=[10, 6])
+fig, ax = plt.subplots()
 
-ax.plot(α_arr, KL_g_arr, label='KL(h, g)')
-ax.plot(α_arr, KL_f_arr, label='KL(h, f)')
-ax.set_ylabel('KL散度')
-ax.set_xlabel(r'$\alpha$')
+ax.plot(α_arr, KL_g_arr, label="KL(h, g)", lw=2)
+ax.plot(α_arr, KL_f_arr, label="KL(h, f)", lw=2)
+ax.set_ylabel("KL divergence")
+ax.set_xlabel(r"$\alpha$")
 
-ax.legend(loc='upper right')
+ax.legend(loc="upper right")
 plt.show()
 ```
 
 让我们计算一个 $\alpha$ 值，使得 $h$ 和 $g$ 之间的 KL 散度等于 $h$ 和 $f$ 之间的 KL 散度。
 
 ```{code-cell} ipython3
-# 当 KL_f = KL_g 时
-discretion = α_arr[np.argmin(np.abs(KL_g_arr-KL_f_arr))]
+# where KL_f = KL_g
+discretion = α_arr[np.argmin(np.abs(KL_g_arr - KL_f_arr))]
 ```
 
 我们可以计算并绘制每个$\alpha$的收敛点$\pi_{\infty}$，以验证收敛确实是由KL散度决定的。
@@ -556,25 +603,35 @@ discretion = α_arr[np.argmin(np.abs(KL_g_arr-KL_f_arr))]
 因此，下图证实了KL散度的最小值如何决定了我们的类型1代理最终学到的内容。
 
 ```{code-cell} ipython3
-α_arr_x = α_arr[(α_arr<discretion)|(α_arr>discretion)]
-π_lim_arr = π_lim_v(α_arr_x)
+---
+mystnb:
+  figure:
+    caption: 极限点与KL散度
+    name: fig-kl-limit
+---
+α_arr_x = α_arr[(α_arr < discretion) | (α_arr > discretion)]
+keys = jax.random.split(jax.random.key(42), len(α_arr_x))
+π_lim_arr = π_lim_v(keys, α_arr_x)
 
 # plot
-fig, ax = plt.subplots(1, figsize=[10, 6])
+fig, ax = plt.subplots()
 
-ax.plot(α_arr, KL_g_arr, label='KL(h, g)')
-ax.plot(α_arr, KL_f_arr, label='KL(h, f)')
-ax.set_ylabel('KL散度')
-ax.set_xlabel(r'$\alpha$')
+ax.plot(α_arr, KL_g_arr, label="KL(h, g)", lw=2)
+ax.plot(α_arr, KL_f_arr, label="KL(h, f)", lw=2)
+ax.set_ylabel("KL divergence")
+ax.set_xlabel(r"$\alpha$")
 
 # plot KL
 ax2 = ax.twinx()
 # plot limit point
-ax2.scatter(α_arr_x, π_lim_arr, 
-            facecolors='none', 
-            edgecolors='tab:blue', 
-            label=r'$\pi$ lim')
-ax2.set_ylabel('π lim')
+ax2.scatter(
+    α_arr_x,
+    π_lim_arr,
+    facecolors="none",
+    edgecolors="tab:blue",
+    label=r"$\pi$ lim",
+)
+ax2.set_ylabel("π lim")
 
 ax.legend(loc=[0.85, 0.8])
 ax2.legend(loc=[0.85, 0.73])
@@ -601,7 +658,7 @@ $$
 
 但不知道$\alpha$。
 
-我们假设此人从$\alpha \in (0,1)$上的先验概率$\pi_0(\alpha)$开始，这个先验具有我们在{doc}`这个quantecon讲座<bayes_nonconj>`中使用的形式之一。
+我们假设此人从$\alpha \in (0,1)$上的先验概率$\pi_0(\alpha)$开始，这个先验具有我们在[](bayes_nonconj)中使用的形式之一。
 
 我们将启动`numpyro`并将其应用于当前情况。
 
@@ -614,9 +671,7 @@ $$
 
 我们将使用numpyro来近似这个方程。
 
-我们将创建后验$\pi_t(\alpha)$的图形，随着
-
-当 $t \rightarrow +\infty$ 时,对应于 quantecon 讲座中展示的内容 <https://python.quantecon.org/bayes_nonconj.html>。
+我们将创建后验$\pi_t(\alpha)$的图形，随着 $t \rightarrow +\infty$，对应于[](bayes_nonconj)中展示的内容。
 
 我们预计后验分布将在 $t \rightarrow + \infty$ 时收敛于真实的 $\alpha$ 值周围。
 
@@ -627,46 +682,65 @@ $$
 ```{code-cell} ipython3
 α = 0.8
 
-# 用真实的 α 模拟数据
-data = draw_lottery(α, 1000)
+# simulate data with true α
+data = draw_lottery(jax.random.key(42), α, 1000)
 sizes = [5, 20, 50, 200, 1000, 25000]
 
+
 def model(w):
-    α = numpyro.sample('α', dist.Uniform(low=0.0, high=1.0))
+    α = numpyro.sample("α", dist.Uniform(low=0.0, high=1.0))
+
+    y_samp = numpyro.sample(
+        "w",
+        dist.Mixture(
+            dist.Categorical(jnp.array([α, 1 - α])),
+            [dist.Beta(F_a, F_b), dist.Beta(G_a, G_b)],
+        ),
+        obs=w,
+    )
 
-    y_samp = numpyro.sample('w',
-        dist.Mixture(dist.Categorical(jnp.array([α, 1-α])), [dist.Beta(F_a, F_b), dist.Beta(G_a, G_b)]), obs=w)
 
 def MCMC_run(ws):
-    "使用 MCMC 计算观测到的 ws 的后验分布"
+    "Compute posterior using MCMC with observed ws"
 
     kernel = NUTS(model)
     mcmc = MCMC(kernel, num_samples=5000, num_warmup=1000, progress_bar=False)
 
-    mcmc.run(rng_key=random.PRNGKey(0), w=jnp.array(ws))
+    mcmc.run(rng_key=random.key(42), w=jnp.array(ws))
     sample = mcmc.get_samples()
-    return sample['α']
+    return sample["α"]
 ```
 
 以下代码生成了下面的图表，显示了不同历史长度下$\alpha$的贝叶斯后验分布。
 
 ```{code-cell} ipython3
-
-fig, ax = plt.subplots(figsize=(10, 6))
+---
+mystnb:
+  figure:
+    caption: 随$t$增长的$\alpha$后验分布
+    name: fig-posterior-alpha
+---
+fig, ax = plt.subplots()
 
 for i in range(len(sizes)):
-    sample = MCMC_run(data[:sizes[i]])
+    sample = MCMC_run(data[: sizes[i]])
     sns.histplot(
-        data=sample, kde=True, stat='density', alpha=0.2, ax=ax,
-        color=colors[i], binwidth=0.02, linewidth=0.05, label=f't={sizes[i]}'
+        data=sample,
+        kde=True,
+        stat="density",
+        alpha=0.2,
+        ax=ax,
+        color=colors[i],
+        binwidth=0.02,
+        linewidth=0.05,
+        label=f"t={sizes[i]}",
     )
-ax.set_title(r'$\pi_t(\alpha)$ as $t$ increases')
 ax.legend()
-ax.set_xlabel(r'$\alpha$')
+ax.set_xlabel(r"$\alpha$")
 plt.show()
 ```
 
-显然，随着观测历史长度的增长，贝叶斯后验分布逐渐收敛于混合参数的真实值 $\alpha = .8$。
+显然，随着观测历史长度的增长，贝叶斯后验分布逐渐收敛于混合参数的真实值 $\alpha = 0.8$。
 
 ## 总结性评论
 
@@ -704,17 +778,17 @@ $s(x | \theta)$来从$x$推断$\theta$。
 ```{exercise}
 :label: mix_model_ex1
 
-在{doc}`likelihood_bayes`中，我们研究了将似然比和贝叶斯定律应用于错误设定统计模型的后果。
+在[](likelihood_bayes)中，我们研究了将似然比和贝叶斯定律应用于错误设定统计模型的后果。
 
 在那节课中，我们使用模型选择算法来研究真实数据生成过程是混合分布的情况。
 
 在本节课中，我们研究了如何使用贝叶斯方法正确地"学习"由混合过程生成的模型。
 
-为了修正我们在{doc}`likelihood_bayes`中使用的算法，正确的贝叶斯方法应该直接对$x$的不确定性建模，并随着新数据的到来更新对它的信念。
+为了修正我们在[](likelihood_bayes)中使用的算法，正确的贝叶斯方法应该直接对$x$的不确定性建模，并随着新数据的到来更新对它的信念。
 
 这是算法：
 
-首先我们指定$x$的先验分布为$x \sim \text{Beta}(\alpha_0, \beta_0)$，其期望值为$\mathbb{E}[x] = \frac{\alpha_0}{\alpha_0 + \beta_0}$。
+首先我们指定$x$的先验分布为$x \sim \mathrm{Beta}(\alpha_0, \beta_0)$，其期望值为$\mathbb{E}[x] = \frac{\alpha_0}{\alpha_0 + \beta_0}$。
 
 单个观测值 $w_t$ 的似然函数为 $p(w_t|x) = x f(w_t) + (1-x) g(w_t)$。
 
@@ -728,7 +802,7 @@ $s(x | \theta)$来从$x$推断$\theta$。
 
 你的任务是用Python实现这个算法。
 
-你可以通过检查后验均值是否随着 $t$ 的增加而收敛到 $x$ 的真实值来验证你的实现，详见 {doc}`likelihood_bayes`。
+你可以通过检查后验均值是否随着 $t$ 的增加而收敛到 $x$ 的真实值来验证你的实现，详见[](likelihood_bayes)。
 ```
 
 ```{solution-start} mix_model_ex1
@@ -746,62 +820,87 @@ T_mix = 200
 
 # 三个不同的先验，均值分别为0.25、0.5、0.75
 prior_params = [(1, 3), (1, 1), (3, 1)]
-prior_means = [a/(a+b) for a, b in prior_params]
+prior_means = [a / (a + b) for a, b in prior_params]
 
-w_mix = draw_lottery(x_true, T_mix)
+w_mix = draw_lottery(jax.random.key(42), x_true, T_mix)
 ```
 
 ```{code-cell} ipython3
-@jit
+@jax.jit
 def learn_x_bayesian(observations, α0, β0, grid_size=2000):
     """
-    使用网格近似对混合概率x进行顺序贝叶斯学习。
+    Sequential Bayesian learning of the mixing probability x
+    using a grid approximation.
     """
-    w = np.asarray(observations)
+    w = jnp.asarray(observations)
     T = w.size
 
-    x_grid = np.linspace(1e-3, 1 - 1e-3, grid_size)
-
-    # 对数先验
-    log_prior = (α0 - 1) * np.log(x_grid) + (β0 - 1) * np.log1p(-x_grid)
-
-    μ_path = np.empty(T + 1)
-    μ_path[0] = α0 / (α0 + β0)
+    x_grid = jnp.linspace(1e-3, 1 - 1e-3, grid_size)
 
-    log_post = log_prior.copy()
+    # Log prior
+    log_prior = (α0 - 1) * jnp.log(x_grid) + (β0 - 1) * jnp.log1p(-x_grid)
 
-    for t in range(T):
-        wt = w[t]
+    def scan_fn(log_post, wt):
         # P(w_t | x) = x f(w_t) + (1 - x) g(w_t)
         like = x_grid * f(wt) + (1 - x_grid) * g(wt)
-        log_post += np.log(like)
+        log_post = log_post + jnp.log(like)
+
+        # Normalize using log-sum-exp trick
+        log_post = log_post - jax.nn.logsumexp(log_post)
+        post = jnp.exp(log_post)
+
+        # Compute posterior mean
+        μ = x_grid @ post
+
+        return log_post, μ
+
+    # Initial posterior mean
+    μ_0 = α0 / (α0 + β0)
 
-        # 归一化
-        log_post -= log_post.max()
-        post = np.exp(log_post)
-        post /= post.sum()
+    # Scan over observations
+    _, μ_path = jax.lax.scan(scan_fn, log_prior, w)
 
-        μ_path[t + 1] = x_grid @ post
+    # Prepend initial value
+    return jnp.concatenate([jnp.array([μ_0]), μ_path])
 
-    return μ_path
 
-x_posterior_means = [learn_x_bayesian(w_mix, α0, β0) for α0, β0 in prior_params]
+# Vectorize over different prior parameters
+def compute_all_posteriors(observations, prior_params):
+    """Compute posterior means for all prior parameter pairs."""
+
+    def single_posterior(params):
+        α0, β0 = params
+        return learn_x_bayesian(observations, α0, β0)
+
+    return jax.vmap(single_posterior)(jnp.array(prior_params))
+
+
+x_posterior_means = compute_all_posteriors(w_mix, jnp.array(prior_params))
 ```
 
 让我们可视化 $x$ 的后验均值如何随时间演变，从三个不同的先验信念开始。
 
 ```{code-cell} ipython3
-fig, ax = plt.subplots(figsize=(10, 6))
+fig, ax = plt.subplots()
 
 for i, (x_means, mean0) in enumerate(zip(x_posterior_means, prior_means)):
-    ax.plot(range(T_mix + 1), x_means, 
-            label=fr'Prior mean = ${mean0:.2f}$', 
-            color=colors[i], linewidth=2)
-
-ax.axhline(y=x_true, color='black', linestyle='--', 
-           label=f'True x = {x_true}', linewidth=2)
-ax.set_xlabel('$t$')
-ax.set_ylabel('Posterior mean of $x$')
+    ax.plot(
+        range(T_mix + 1),
+        x_means,
+        label=rf"Prior mean = ${mean0:.2f}$",
+        color=colors[i],
+        linewidth=2,
+    )
+
+ax.axhline(
+    y=x_true,
+    color="black",
+    linestyle="--",
+    label=f"True x = {x_true}",
+    linewidth=2,
+)
+ax.set_xlabel("$t$")
+ax.set_ylabel("Posterior mean of $x$")
 ax.legend()
 plt.show()
 ```
@@ -811,21 +910,27 @@ plt.show()
 接下来，让我们看看更长时间范围内的多次模拟，所有模拟都从均匀先验开始。
 
 ```{code-cell} ipython3
-set_seed()
 n_paths = 20
 T_long = 10_000
 
-fig, ax = plt.subplots(figsize=(10, 5))
+keys = jax.random.split(jax.random.key(42), n_paths)
+
+fig, ax = plt.subplots()
 
 for j in range(n_paths):
-    w_path = draw_lottery(x_true, T_long) 
-    x_means = learn_x_bayesian(w_path, 1, 1)  # 均匀先验
-    ax.plot(range(T_long + 1), x_means, alpha=0.5, linewidth=1)
-
-ax.axhline(y=x_true, color='red', linestyle='--', 
-            label=f'True x = {x_true}', linewidth=2)
-ax.set_ylabel('Posterior mean of $x$')
-ax.set_xlabel('$t$')
+    w_path = draw_lottery(keys[j], x_true, T_long)
+    x_means = learn_x_bayesian(w_path, 1, 1)  # Uniform prior
+    ax.plot(range(T_long + 1), x_means, alpha=0.5, lw=2)
+
+ax.axhline(
+    y=x_true,
+    color="red",
+    linestyle="--",
+    label=f"True x = {x_true}",
+    linewidth=2,
+)
+ax.set_ylabel("Posterior mean of $x$")
+ax.set_xlabel("$t$")
 ax.legend()
 plt.tight_layout()
 plt.show()
@@ -834,5 +939,4 @@ plt.show()
 我们可以看到 $x$ 的后验均值收敛到真实值 $x=0.5$。
 
 ```{solution-end}
-```
-
+```
\ No newline at end of file