第 7 章　实战专题

# 公共导入与工具（后面所有单元都依赖这里）
import time, math, logging, warnings, functools, dataclasses
from dataclasses import dataclass, field
from functools import partial
from typing import NamedTuple, Any, Callable

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.monitoring as monitoring
import matplotlib.pyplot as plt

plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False
# 图表配色（固定顺序的分类色）
C_BLUE, C_ORANGE, C_AQUA, C_VIOLET, C_RED = "#2a78d6", "#eb6834", "#1baf7a", "#6250d6", "#e34948"

print("JAX", jax.__version__, "| 后端", jax.default_backend(), "| 设备", jax.devices())


def bench(f, *args, warmup=2, runs=10):
    """计时：先预热 warmup 次（含编译），再运行 runs 次取中位数，返回毫秒。"""
    for _ in range(warmup):
        jax.block_until_ready(f(*args))
    ts = []
    for _ in range(runs):
        t0 = time.perf_counter()
        jax.block_until_ready(f(*args))
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts)) * 1e3


def assert_faster(fast, slow, factor=1.0, tries=3):
    """计时断言：fast() 的耗时 × factor < slow() 的耗时（两个参数都是返回耗时的无参函数）。
    计时受机器负载影响（例如 CI 上几个 notebook 并行），不满足时重新测量，最多 tries 次。"""
    for _ in range(tries):
        tf, ts = fast(), slow()
        if tf * factor < ts:
            return tf, ts
    raise AssertionError(f"重测 {tries} 次仍不满足：{tf:.3g} × {factor} ≥ {ts:.3g}")

def show_error(fn):
    """运行 fn；若报错，只打印异常类型和第一行信息（用于演示错误写法）。"""
    try:
        out = fn()
        print("（没有报错）结果:", out)
    except Exception as e:
        first = str(e).strip().splitlines()[0] if str(e).strip() else ""
        print(f"❌ {type(e).__name__}: {first[:160]}")


_COMPILE_EVENT = "/jax/core/compile/backend_compile_duration"


class CompileCounter:
    """统计 with 块内发生的 XLA 后端编译次数。

    name=None 统计所有编译；name="f" 只统计 jit(f) 的编译。
    原理：JAX 每完成一次后端编译，都会通过 jax.monitoring 发出一个
    '/jax/core/compile/backend_compile_duration' 事件，附带 fun_name。
    """
    def __init__(self, name=None):
        self.name = name
        self.events = []                      # [(fun_name, 秒)]
        self._listener = self._on_event       # 保存同一个绑定方法对象，注销时才能找到它

    def _on_event(self, event, duration, **kw):
        if event != _COMPILE_EVENT:
            return
        fun = kw.get("fun_name", "?")
        if self.name is None or fun == f"jit({self.name})":
            self.events.append((fun, duration))

    def __enter__(self):
        monitoring.register_event_duration_secs_listener(self._listener)
        return self

    def __exit__(self, *exc):
        monitoring.unregister_event_duration_listener(self._listener)
        return False

    @property
    def count(self):
        return len(self.events)


def wpad(s, width):
    """按显示宽度左对齐（一个中文字符占 2 列），用于打印对齐的表格。"""
    import unicodedata
    w = sum(2 if unicodedata.east_asian_width(ch) in "WF" else 1 for ch in str(s))
    return str(s) + " " * max(0, width - w)


def to64(x):
    """转成 NumPy float64，用于和高精度参考值对拍。"""
    return np.asarray(x, dtype=np.float64)

# 自检：计数器能数到编译，命中缓存时不计数
_x = jnp.ones(3)
def _probe(x):
    return x * 2
with CompileCounter("_probe") as _c:
    jax.jit(_probe)(_x); jax.jit(_probe)(_x)
assert _c.count == 1, _c.events
print("CompileCounter 自检通过：两次调用只编译 1 次")

def mlp_loss(w1, w2, x):
    # 一个小 MLP 的前向 + 均方值，权重和数据都作为参数传入
    h = jnp.tanh(x @ w1)
    return jnp.mean((h @ w2) ** 2)

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
W1 = jax.random.normal(k1, (64, 256)) / 8.0
W2 = jax.random.normal(k2, (256, 10)) / 16.0
X_all = jax.random.normal(k3, (600, 64))
grad_fn = jax.jit(jax.grad(mlp_loss, argnums=(0, 1)))

sizes = [100 + 7 * i for i in range(20)]            # 20 种不同的 batch 大小
batches = [X_all[:n] for n in sizes]                # 先切好，避免把切片耗时算进去
same = [X_all[:100]] * 20

def run(xs):
    t0 = time.perf_counter()
    for x in xs:
        jax.block_until_ready(grad_fn(W1, W2, x))
    return (time.perf_counter() - t0) * 1e3

with CompileCounter("mlp_loss") as c_var:
    t_var = run(batches)
grad_fn(W1, W2, same[0])                             # 预热：保证下面只测执行
with CompileCounter("mlp_loss") as c_same:
    t_same = run(same)
print(f"20 种 batch 大小：{t_var:7.1f} ms，编译 {c_var.count} 次")
print(f"同一 batch 大小 ：{t_same:7.1f} ms，编译 {c_same.count} 次")
print(f"平均每次编译约 {np.mean([d for _, d in c_var.events]) * 1e3:.1f} ms，"
      f"每次执行约 {t_same / 20:.2f} ms")
assert c_var.count == 20 and c_same.count == 0

x3, x4 = jnp.ones(3), jnp.ones(4)
x3_i = jnp.ones(3, jnp.int32)

def count_compiles(name, calls):
    """依次执行 calls 中的每个无参函数，返回 jit(name) 被编译的次数。"""
    with CompileCounter(name) as c:
        for call in calls:
            jax.block_until_ready(call())
    return c.count

rows = []   # (触发因素, 调用方式, 实测编译次数, 预期次数)

# ① 形状
def t_shape(x): return x * 2
f = jax.jit(t_shape)
rows.append(("形状", "f32[3] → f32[4]", count_compiles("t_shape", [lambda: f(x3), lambda: f(x4)]), 2))

# ② dtype
def t_dtype(x): return x * 2
f = jax.jit(t_dtype)
rows.append(("dtype", "f32[3] → i32[3]", count_compiles("t_dtype", [lambda: f(x3), lambda: f(x3_i)]), 2))

# ③ 弱类型：Python 标量是“弱类型”，数组是“强类型”
def t_weak(x, lr): return x * lr
f = jax.jit(t_weak)
rows.append(("弱类型", "lr=0.1 → lr=jnp.float32(0.1)",
             count_compiles("t_weak", [lambda: f(x3, 0.1), lambda: f(x3, jnp.float32(0.1))]), 2))

# ④ 静态参数的值
def t_static(x, n): return x ** n
f = jax.jit(t_static, static_argnames="n")
rows.append(("静态参数值", "n=2 → n=3", count_compiles("t_static", [lambda: f(x3, n=2), lambda: f(x3, n=3)]), 2))

# ⑤ pytree 结构：多一个键 / list 换成 tuple / 叶子换成 None
def t_tree(s): return jax.tree.map(lambda a: a + 1, s)
f = jax.jit(t_tree)
rows.append(("pytree 结构", "{'a'} → {'a','b'} → [x] → (x,) → {'a', 'b':None}",
             count_compiles("t_tree", [lambda: f({"a": x3}), lambda: f({"a": x3, "b": x3}),
                                       lambda: f([x3]), lambda: f((x3,)), lambda: f({"a": x3, "b": None})]), 5))

# ⑥ 新建 jit 对象：每次调用都创建一个新的内部函数（闭包）
def make_and_call(x):
    def t_inner(y): return y + 1          # 每次调用 make_and_call 都是一个新的函数对象
    return jax.jit(t_inner)(x)
rows.append(("新建 jit 对象（闭包）", "循环里调用 3 次",
             count_compiles("t_inner", [lambda: make_and_call(x3)] * 3), 3))

# ⑦ 新建 jit 对象：jit 一个绑定方法（obj.method 每次访问都生成新的绑定方法对象）
class Model:
    def t_method(self, x): return x * 3
    def predict(self, x): return jax.jit(self.t_method)(x)
m = Model()
rows.append(("新建 jit 对象（绑定方法）", "m.predict 调用 3 次",
             count_compiles("t_method", [lambda: m.predict(x3)] * 3), 3))

# ⑧ 静态参数按 id 比较：自定义类没有 __eq__/__hash__，每个新实例都是新键
class Cfg:
    def __init__(self, n): self.n = n
def t_cfg(x, cfg): return x * cfg.n
f = jax.jit(t_cfg, static_argnums=1)
rows.append(("静态参数按 id 哈希", "Cfg(2) 新建 3 次", count_compiles("t_cfg", [lambda: f(x3, Cfg(2))] * 3), 3))

# ⑨ 调用方式：省略默认参数 vs 显式传入相同的值
def t_default(x, scale=2.0): return x * scale
f = jax.jit(t_default)
rows.append(("调用方式", "f(x) → f(x, 2.0) → f(x, scale=2.0)",
             count_compiles("t_default", [lambda: f(x3), lambda: f(x3, 2.0), lambda: f(x3, scale=2.0)]), 3))

# —— 对照组：这些不会触发重新编译 ——
def t_value(x): return x * 2
f = jax.jit(t_value)
rows.append(("（对照）只改数值", "ones(3) → ones(3)*5", count_compiles("t_value", [lambda: f(x3), lambda: f(x3 * 5)]), 1))

def t_shared(x): return x * 2
rows.append(("（对照）多次 jax.jit(f)", "同一个模块级函数包装 3 次",
             count_compiles("t_shared", [lambda: jax.jit(t_shared)(x3)] * 3), 1))

print(wpad("触发因素", 26) + wpad("调用", 50) + "实测  预期")
print("-" * 80)
for name, how, got, exp in rows:
    print(wpad(name, 26) + wpad(how, 50) + f"{got:>4}{exp:>6}")
for name, how, got, exp in rows:
    assert got == exp, (name, got, exp)
print("\n全部与预期一致 ✔")

class capture_jax_logs:
    """临时接管 'jax' logger，把日志收集到 self.messages（不再打印到 stderr）。"""
    def __enter__(self):
        self.logger = logging.getLogger("jax")
        self.messages = []
        self.handler = logging.Handler()
        self.handler.emit = lambda record: self.messages.append(record.getMessage())
        self._old_propagate = self.logger.propagate
        self.logger.addHandler(self.handler)
        self.logger.propagate = False
        return self

    def __exit__(self, *exc):
        self.logger.removeHandler(self.handler)
        self.logger.propagate = self._old_propagate
        return False

# ① jax_log_compiles：每次真正编译打印一行 “Compiling jit(...)”
def diag_step(params, x):
    return jax.tree.map(lambda p: p * x.mean(), params)

f = jax.jit(diag_step)
params_a = {"w": jnp.ones((2, 2)), "b": jnp.zeros(2)}
xs = [jnp.ones(5), jnp.ones(5) * 2, jnp.ones(6)]           # 先建好输入，免得混入 jnp.ones 自己的编译日志

jax.config.update("jax_log_compiles", True)
try:
    with capture_jax_logs() as logs:
        for x in xs:
            f(params_a, x)
finally:
    jax.config.update("jax_log_compiles", False)            # 用 finally 保证一定关掉
compiling = [m for m in logs.messages if m.startswith("Compiling")]
for m in compiling:
    print(m[:150], "...")
assert len(compiling) == 2          # 第 2 次只是数值不同，命中缓存

# ② jax_explain_cache_misses：直接告诉你“哪个参数、哪里变了”
def diag_explain(state, lr):
    return jax.tree.map(lambda p: p - lr * p, state)

f = jax.jit(diag_explain)
s1 = {"w": jnp.ones(3)}
s2 = {"w": jnp.ones(3), "m": jnp.zeros(3)}                  # 结构变了
f(s1, 0.1)                                                  # 先建立缓存
jax.config.update("jax_explain_cache_misses", True)
try:
    with capture_jax_logs() as logs:
        f(s1, jnp.float32(0.1))        # 弱类型 → 强类型
        f(s2, 0.1)                     # pytree 结构变化
finally:
    jax.config.update("jax_explain_cache_misses", False)
misses = [m for m in logs.messages if "CACHE MISS" in m]
for m in misses:
    # 去掉文件路径行，只保留“原因”部分
    lines = [l for l in m.splitlines() if not l.strip().startswith(("for ", "TRACING"))]
    print("\n".join(lines[:6]), "\n")
assert len(misses) == 2

trace_count = 0
def diag_counter(x):
    global trace_count
    trace_count += 1               # 只在追踪时执行
    return jnp.sin(x) * 2

f = jax.jit(diag_counter)
with CompileCounter("diag_counter") as c:
    jax.eval_shape(f, x3)          # 只追踪，算出输出的形状/dtype
    f.lower(x4)                    # 追踪 + 降级到 StableHLO，不编译
    f(x3); f(x3)                   # 第一次：追踪 + 编译（复用 eval_shape 的追踪结果）；第二次：命中
print(f"追踪次数 = {trace_count}，后端编译次数 = {c.count}")
assert c.count == 1 and trace_count >= 2

def deep_unrolled(ws, x):
    for w in ws:                   # Python 循环在追踪时展开成 30 份
        x = jnp.tanh(x @ w)
    return x.sum()

ws = [jax.random.normal(k, (64, 64)) / 8.0 for k in jax.random.split(jax.random.key(1), 30)]
xb = jnp.ones((16, 64))
jf = jax.jit(deep_unrolled)

t0 = time.perf_counter(); traced = jf.trace(ws, xb); t1 = time.perf_counter()
lowered = traced.lower();                                t2 = time.perf_counter()
compiled = lowered.compile();                            t3 = time.perf_counter()
print(f"trace {1e3*(t1-t0):6.1f} ms | lower {1e3*(t2-t1):6.1f} ms | compile {1e3*(t3-t2):6.1f} ms")
print("jaxpr 方程数:", len(traced.jaxpr.eqns), "| StableHLO 行数:", len(lowered.as_text().splitlines()))
print("类型:", type(traced).__name__, type(lowered).__name__, type(compiled).__name__,
      "| 都在 jax.stages 里:", all(hasattr(jax.stages, n) for n in ("Traced", "Lowered", "Compiled")))

ref = to64(xb)                                            # NumPy float64 参考值
for w in ws:
    ref = np.tanh(ref @ to64(w))
assert np.allclose(compiled(ws, xb), ref.sum(), rtol=1e-4, atol=1e-4)
show_error(lambda: compiled(ws, jnp.ones((17, 64))))       # 形状不同 → TypeError

def masked_mse(w, X, y, mask):
    r = X @ w - y
    return jnp.sum(mask * r ** 2) / jnp.sum(mask)          # 分子、分母都只算有效行

def plain_mse(w, X, y):
    return jnp.mean((X @ w - y) ** 2)

BUCKETS = (64, 128, 256, 512)

def pad_to_bucket(X, y):
    """在 NumPy 里做 padding（不触发任何 JAX 编译），返回 (Xp, yp, mask)。"""
    n = X.shape[0]
    b = next(b for b in BUCKETS if b >= n)
    Xp = np.zeros((b, X.shape[1]), np.float32); Xp[:n] = X
    yp = np.zeros(b, np.float32); yp[:n] = y
    mask = np.zeros(b, np.float32); mask[:n] = 1.0
    return Xp, yp, mask

kx, ky, kw, kn = jax.random.split(jax.random.key(2), 4)
Xbig = np.asarray(jax.random.normal(kx, (512, 8)))
ybig = np.asarray(jax.random.normal(ky, (512,)))
w0 = jax.random.normal(kw, (8,))
lengths = [int(n) for n in jax.random.randint(kn, (30,), 20, 500)]   # 30 个随机长度

vg_plain = jax.jit(jax.value_and_grad(plain_mse))
vg_masked = jax.jit(jax.value_and_grad(masked_mse))
wrong_mean = jax.jit(lambda w, X, y: jnp.mean((X @ w - y) ** 2))      # ❌ padding 了但没 mask

with CompileCounter("plain_mse") as c_plain:
    for n in lengths:
        vg_plain(w0, Xbig[:n], ybig[:n])
max_err_masked, max_err_wrong = 0.0, 0.0
with CompileCounter("masked_mse") as c_masked:
    for n in lengths:
        Xp, yp, mask = pad_to_bucket(Xbig[:n], ybig[:n])
        loss, g = vg_masked(w0, Xp, yp, mask)
        # NumPy float64 参考值
        r = Xbig[:n].astype(np.float64) @ to64(w0) - ybig[:n]
        ref_loss, ref_g = np.mean(r ** 2), 2 * Xbig[:n].T.astype(np.float64) @ r / n
        max_err_masked = max(max_err_masked, abs(float(loss) - ref_loss), np.abs(to64(g) - ref_g).max())
        max_err_wrong = max(max_err_wrong, abs(float(wrong_mean(w0, Xp, yp)) - ref_loss))

print(f"不分桶：{len(set(lengths))} 种长度 → 编译 {c_plain.count} 次")
print(f"分桶  ：{len(BUCKETS)} 个桶     → 编译 {c_masked.count} 次")
print(f"分桶 + mask 与 NumPy float64 的最大误差：{max_err_masked:.2e}")
print(f"❌ padding 但不 mask 的损失最大误差    ：{max_err_wrong:.3f}")
assert c_plain.count == len(set(lengths)) and c_masked.count <= len(BUCKETS)
assert max_err_masked < 1e-4 and max_err_wrong > 0.1

# padding 的代价：平均浪费了多少计算
waste = np.mean([1 - n / next(b for b in BUCKETS if b >= n) for n in lengths])
print(f"2 的幂分桶平均浪费 {waste:.0%} 的计算；桶分得越细，浪费越少，但编译次数越多")

@dataclass(frozen=True)          # frozen=True → 自动生成按值比较的 __eq__ 和 __hash__
class CfgOK:
    n: int

def t_cfg_ok(x, cfg): return x * cfg.n
f = jax.jit(t_cfg_ok, static_argnums=1)
n_ok = count_compiles("t_cfg_ok", [lambda: f(x3, CfgOK(2))] * 3)

@jax.tree_util.register_static   # 注册为静态 pytree：作为普通参数传入即可
@dataclass(frozen=True)
class Activation:
    name: str

def t_act(x, act):
    return jnp.tanh(x) if act.name == "tanh" else jax.nn.relu(x)   # Python if：act 是静态的
f_act = jax.jit(t_act)           # 注意：没有写 static_argnums
n_act = count_compiles("t_act", [lambda: f_act(x3, Activation("tanh"))] * 3 + [lambda: f_act(x3, Activation("relu"))])

# 学习率：❌ 静态（每个新值编译一次） vs ✅ 动态（只编译一次）
def t_lr_static(p, lr): return p - lr * p
def t_lr_dynamic(p, lr): return p - lr * p
f_s = jax.jit(t_lr_static, static_argnames="lr")
f_d = jax.jit(t_lr_dynamic)
lrs = [1e-3 * (1 + i) for i in range(8)]                      # 模拟 8 步学习率调度
n_s = count_compiles("t_lr_static", [lambda lr=lr: f_s(x3, lr=lr) for lr in lrs])
n_d = count_compiles("t_lr_dynamic", [lambda lr=lr: f_d(x3, jnp.float32(lr)) for lr in lrs])

print(f"frozen dataclass 作静态参数，3 个等值实例 → 编译 {n_ok} 次")
print(f"register_static：tanh × 3 + relu × 1     → 编译 {n_act} 次")
print(f"学习率调度 8 步：静态 {n_s} 次 | 动态 {n_d} 次")
assert (n_ok, n_act, n_s, n_d) == (1, 2, 8, 1)

def _forward(w, x):                         # 纯函数：权重显式传入
    return jnp.tanh(x @ w)
forward_jit = jax.jit(_forward)             # ✅ 模块级创建一次

class ModelOK:
    def __init__(self, w):
        self.w = w
    def predict(self, x):
        return forward_jit(self.w, x)       # 方法里只调用，不新建 jit

mo = ModelOK(jnp.ones((3, 2)))
n_fwd = count_compiles("_forward", [lambda: mo.predict(x3)] * 5)

def t_lr_mixed(p, lr): return p - lr * p
f_mix = jax.jit(t_lr_mixed)
schedule = [0.1, jnp.float32(0.05), 0.02]                     # ❌ Python float 和数组混用
n_mix = count_compiles("t_lr_mixed", [lambda lr=lr: f_mix(x3, lr) for lr in schedule])
def t_lr_cast(p, lr): return p - lr * p
f_cast = jax.jit(t_lr_cast)
n_cast = count_compiles("t_lr_cast", [lambda lr=lr: f_cast(x3, jnp.asarray(lr, jnp.float32)) for lr in schedule])  # ✅
def t_lr_py(p, lr): return p - lr * p
f_py = jax.jit(t_lr_py)
n_py = count_compiles("t_lr_py", [lambda lr=lr: f_py(x3, lr) for lr in [0.1, 0.05, 0.02]])  # 全是 Python float：值变了也不重编译

print(f"模块级 jit，predict 调用 5 次 → 编译 {n_fwd} 次")
print(f"学习率混用 Python float / 数组 → 编译 {n_mix} 次；统一转成 float32 数组 → 编译 {n_cast} 次；全用 Python float → 编译 {n_py} 次")
assert (n_fwd, n_mix, n_cast, n_py) == (1, 2, 1, 1)

def serve_loss(w, X, y, mask):
    return masked_mse(w, X, y, mask)

t0 = time.perf_counter()
compiled_by_bucket = {
    b: jax.jit(serve_loss).lower(
        jax.ShapeDtypeStruct((8,), jnp.float32), jax.ShapeDtypeStruct((b, 8), jnp.float32),
        jax.ShapeDtypeStruct((b,), jnp.float32), jax.ShapeDtypeStruct((b,), jnp.float32)).compile()
    for b in BUCKETS}
print(f"启动时预编译 {len(BUCKETS)} 个桶：{(time.perf_counter() - t0) * 1e3:.0f} ms")

w0_np = np.asarray(w0)
with CompileCounter() as c_serve:                 # name=None：统计所有函数的编译
    outs = []
    for n in lengths:
        Xp, yp, mask = pad_to_bucket(Xbig[:n], ybig[:n])      # NumPy 完成 padding
        outs.append(compiled_by_bucket[Xp.shape[0]](w0_np, Xp, yp, mask))
    jax.block_until_ready(outs)
print(f"服务 {len(lengths)} 个请求期间发生的编译：{c_serve.count} 次")
assert c_serve.count == 0

class CompileBudgetExceeded(RuntimeError):
    pass


class CompileBudget(CompileCounter):
    """编译预算：超过 max_compiles 次就报错（或警告）。"""
    def __init__(self, max_compiles, name=None, on_exceed="raise"):
        super().__init__(name)
        self.max_compiles = max_compiles
        self.on_exceed = on_exceed

    def check(self):
        if self.count > self.max_compiles:
            detail = ", ".join(f"{fn}({d * 1e3:.0f}ms)" for fn, d in self.events)
            msg = (f"编译预算超支：{self.name or '全部函数'} 允许 {self.max_compiles} 次，"
                   f"实际 {self.count} 次 → {detail}")
            if self.on_exceed == "raise":
                raise CompileBudgetExceeded(msg)
            warnings.warn(msg, RuntimeWarning, stacklevel=3)

    def __exit__(self, exc_type, exc, tb):
        super().__exit__(exc_type, exc, tb)
        if exc_type is None:          # with 块本身出错时不再叠加预算错误
            self.check()
        return False


def compile_budget(max_compiles, on_exceed="raise"):
    """装饰 jit 函数：累计该函数的编译次数，每次调用后检查预算。"""
    def deco(jitted):
        budget = CompileBudget(max_compiles, jitted.__name__, on_exceed)

        @functools.wraps(jitted)
        def wrapper(*args, **kwargs):
            monitoring.register_event_duration_secs_listener(budget._listener)
            try:
                out = jitted(*args, **kwargs)
            finally:
                monitoring.unregister_event_duration_listener(budget._listener)
            budget.check()
            return out

        wrapper.budget = budget
        return wrapper
    return deco

# 用法 1：装饰器。允许 2 次编译（比如训练和评估各一种 batch 形状）
@compile_budget(2)
@jax.jit
def budget_step(w, x):
    return w - 0.1 * jnp.mean(x) * w

w_ = jnp.ones(4)
for n in [32, 32, 64, 64, 32]:                 # 只有 2 种形状：在预算内
    w_ = budget_step(w_, jnp.ones(n))
print("5 次调用、2 种形状：编译", budget_step.budget.count, "次，未超预算")
assert budget_step.budget.count == 2

try:
    budget_step(w_, jnp.ones(48))               # 第 3 种形状 → 超预算
    raise AssertionError("应该报错")
except CompileBudgetExceeded as e:
    print("❌ 捕获到:", str(e)[:120])

# 用法 2：上下文管理器 + warn 模式
def budget_loop(p, lr): return p - lr * p
f = jax.jit(budget_loop)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    with CompileBudget(1, "budget_loop", on_exceed="warn") as b:
        for lr in [0.1, 0.05, jnp.float32(0.01)]:   # 混入一个强类型 → 多编译一次
            f(x3, lr)
print("⚠️ 警告:", str(caught[0].message)[:110])
assert b.count == 2 and len(caught) == 1

# 用法 3：✅ 写成断言——修好之后在预算内，安静通过
with CompileBudget(1, "budget_loop"):
    for lr in [0.1, 0.05, 0.01]:
        f(x3, jnp.asarray(lr, jnp.float32))         # 强类型 f32[] 已在上面编译过 → 0 次
print("修复后：在预算内 ✔")

# 一个 NaN 样本让整组参数一步报废
def per_sample_loss(w, x):
    return jnp.sqrt(jnp.sum((w * x) ** 2))      # ‖w ⊙ x‖：x 全为 0 时梯度是 0/0

w_nan = jnp.ones(4)
X_nan = jnp.stack([jnp.ones(4), 2 * jnp.ones(4), jnp.zeros(4)])     # 第 3 个样本全为 0
g = jax.grad(lambda w: jnp.mean(jax.vmap(per_sample_loss, (None, 0))(w, X_nan)))(w_nan)
w_after = w_nan - 0.1 * g
print("float32 exp 的边界：exp(88.7) =", jnp.exp(88.7), "| exp(89) =", jnp.exp(89.0),
      "| exp(-87) =", jnp.exp(-87.0), "| exp(-88) =", jnp.exp(-88.0), "| NumPy exp(-100) =", np.exp(np.float32(-100.0)))
assert float(jnp.exp(-88.0)) == 0.0 and float(jnp.exp(-87.0)) > 0 and np.exp(np.float32(-100.0)) > 0   # JAX 冲零，NumPy 保留非规格化数
print("batch 梯度:", g)
print("更新一步后的参数:", w_after)
assert bool(jnp.all(jnp.isnan(w_after)))

def double_where(fn, x, ok, safe_value, fill):
    """在 ok 为 False 的位置返回 fill，且保证那里的梯度是有限值。"""
    x_safe = jnp.where(ok, x, safe_value)
    return jnp.where(ok, fn(x_safe), fill)

# ---- 4 个危险函数：❌ 单 where 版 与 ✅ 双 where 版 ----
def norm_bad(v):                                   # 向量范数，零向量处 0/0
    return jnp.linalg.norm(v)
def norm_good(v):
    sq = jnp.sum(v ** 2)
    return double_where(jnp.sqrt, sq, sq > 0, 1.0, 0.0)

def sinc_bad(x):                                   # sin(x)/x，x=0 处补 1
    return jnp.where(x != 0, jnp.sin(x) / x, 1.0)
def sinc_good(x):
    return double_where(lambda t: jnp.sin(t) / t, x, x != 0, 1.0, 1.0)

def entropy_bad(p):                                # -Σ p log p，p=0 处约定为 0
    return -jnp.sum(jnp.where(p > 0, p * jnp.log(p), 0.0))
def entropy_good(p):
    return -jnp.sum(double_where(lambda t: t * jnp.log(t), p, p > 0, 1.0, 0.0))

def logit_bad(p):                                  # log(p/(1-p))，p∈{0,1} 时截断到 ±20
    return jnp.where((p > 0) & (p < 1), jnp.log(p / (1 - p)), 20.0 * jnp.sign(p - 0.5))
def logit_good(p):
    ok = (p > 0) & (p < 1)
    return double_where(lambda t: jnp.log(t / (1 - t)), p, ok, 0.5, 20.0 * jnp.sign(p - 0.5))

cases = [
    ("范数 ‖v‖ @ v=0",        norm_bad,    norm_good,    jnp.zeros(3)),
    ("sinc @ x=0",            sinc_bad,    sinc_good,    jnp.array(0.0)),
    ("熵 @ p=[0, .5, .5]",    entropy_bad, entropy_good, jnp.array([0.0, 0.5, 0.5])),
    ("logit @ p=1",           logit_bad,   logit_good,   jnp.array(1.0)),
]
print(wpad("函数 @ 危险点", 24) + wpad("前向(❌/✅)", 22) + "梯度 ❌  →  梯度 ✅")
for name, bad, good, x in cases:
    gb, gg = jax.grad(bad)(x), jax.grad(good)(x)
    print(wpad(name, 24) + wpad(f"{float(bad(x)):.3f} / {float(good(x)):.3f}", 22) + f"{gb}  →  {gg}")
    assert np.allclose(bad(x), good(x))                       # 前向值完全相同
    assert not bool(jnp.all(jnp.isfinite(gb)))                # ❌ 梯度里有 NaN
    assert bool(jnp.all(jnp.isfinite(gg)))                    # ✅ 梯度全部有限

# 在正常点上，✅ 版本的梯度必须等于解析导数（NumPy float64 计算）
v, x, p, q = np.array([3.0, 4.0, 12.0]), 0.7, np.array([0.2, 0.3, 0.5]), 0.3
assert np.allclose(jax.grad(norm_good)(jnp.array(v)), v / np.linalg.norm(v), atol=1e-6)
assert np.allclose(jax.grad(sinc_good)(x), (x * np.cos(x) - np.sin(x)) / x ** 2, atol=1e-6)
assert np.allclose(jax.grad(entropy_good)(jnp.array(p)), -(np.log(p) + 1), atol=1e-6)
assert np.allclose(jax.grad(logit_good)(q), 1 / (q * (1 - q)), atol=1e-4)
print("\n正常点上的梯度与解析导数一致 ✔")

def logsumexp_stable(x, axis=-1):
    m = jnp.max(x, axis=axis, keepdims=True)
    m = lax.stop_gradient(jnp.where(jnp.isfinite(m), m, 0.0))     # 细节 1、2
    return jnp.squeeze(m, axis) + jnp.log(jnp.sum(jnp.exp(x - m), axis=axis))

def log_softmax_stable(x, axis=-1):
    return x - jnp.expand_dims(logsumexp_stable(x, axis), axis)

def softmax_stable(x, axis=-1):
    m = lax.stop_gradient(jnp.max(x, axis=axis, keepdims=True))
    e = jnp.exp(x - m)
    return e / jnp.sum(e, axis=axis, keepdims=True)

# ❌ 朴素写法
def logsumexp_naive(x): return jnp.log(jnp.sum(jnp.exp(x), axis=-1))
def softmax_naive(x):   return jnp.exp(x) / jnp.sum(jnp.exp(x), axis=-1, keepdims=True)
def log_softmax_naive(x): return jnp.log(softmax_naive(x))

def np_log_softmax(x):          # NumPy float64 参考
    x = np.asarray(x, np.float64)
    m = x.max(-1, keepdims=True)
    return x - (m + np.log(np.exp(x - m).sum(-1, keepdims=True)))

logits = jnp.array([[1000.0, 0.0, -1000.0],     # 上溢
                    [-1000.0, -1001.0, -1002.0], # 全部下溢
                    [0.0, -120.0, 3.0],          # 有一个概率极小
                    [1.0, 2.0, 3.0]])            # 普通情况
ref = np_log_softmax(logits)
print("朴素 log(softmax):\n", log_softmax_naive(logits))
print("稳定 log_softmax:\n", log_softmax_stable(logits))
print("朴素 logsumexp:", logsumexp_naive(logits), "\n稳定 logsumexp:", logsumexp_stable(logits))

assert np.allclose(log_softmax_stable(logits), ref, atol=1e-4)
assert np.allclose(log_softmax_stable(logits), jax.nn.log_softmax(logits), atol=1e-5)
assert np.allclose(softmax_stable(logits), np.exp(ref), atol=1e-6)
assert np.allclose(logsumexp_stable(logits), jax.nn.logsumexp(logits, axis=-1), atol=1e-4)
assert not bool(jnp.all(jnp.isfinite(log_softmax_naive(logits))))
print("稳定版与 NumPy float64、jax.nn 一致 ✔")

def ce_naive(z, y):
    return -jnp.mean(jnp.take_along_axis(log_softmax_naive(z), y[:, None], axis=1))

def ce_stable(z, y):
    return -jnp.mean(jnp.take_along_axis(log_softmax_stable(z), y[:, None], axis=1))

z = jnp.array([[0.0, -120.0, 3.0], [1.0, 2.0, 3.0]])
y = jnp.array([1, 2])                        # 第 1 个样本的正确类别概率约为 e^-123
for name, fn in [("❌ log(softmax)", ce_naive), ("✅ log_softmax", ce_stable)]:
    loss, gz = jax.value_and_grad(fn)(z, y)
    print(f"{name}: loss = {float(loss):.4f}，梯度有限 = {bool(jnp.all(jnp.isfinite(gz)))}")

loss, gz = jax.value_and_grad(ce_stable)(z, y)
ref_g = (np.exp(np_log_softmax(z)) - np.eye(3)[np.asarray(y)]) / len(y)       # (softmax - onehot) / B
ref_loss = -np.mean(np_log_softmax(z)[np.arange(2), np.asarray(y)])
assert np.allclose(gz, ref_g, atol=1e-6) and np.isclose(float(loss), ref_loss, rtol=1e-5)
assert not np.isfinite(float(ce_naive(z, y)))
print("梯度 = (softmax − onehot)/B，与 NumPy float64 一致 ✔")

NEG = jnp.finfo(jnp.float32).min

def masked_ce(z, y, valid_cols, row_weight, fill):
    z = jnp.where(valid_cols, z, fill)                       # 屏蔽无效类别
    logp = jax.nn.log_softmax(z)
    nll = -jnp.take_along_axis(logp, y[:, None], axis=1)[:, 0]
    return jnp.sum(nll * row_weight) / jnp.sum(row_weight), jax.nn.softmax(z)

zb = jnp.array([[2.0, 1.0, 0.5], [0.3, 0.2, 0.1]])
yb = jnp.array([0, 0])
valid_cols = jnp.array([[True, True, False], [False, False, False]])   # 第 2 行是 padding：全被 mask
row_weight = jnp.array([1.0, 0.0])                                      # 浮点权重：padding 行为 0

out = {}
for name, fill in [("❌ 填 -inf", -jnp.inf), ("✅ 填 finfo.min", NEG)]:
    (loss, probs), gz = jax.value_and_grad(masked_ce, has_aux=True)(zb, yb, valid_cols, row_weight, fill)
    out[name] = (loss, probs, gz)
    print(f"{name}: loss = {float(loss)}，梯度有限 = {bool(jnp.all(jnp.isfinite(gz)))}，"
          f"padding 行的概率 = {probs[1]}")

loss_ok, probs_ok, g_ok = out["✅ 填 finfo.min"]
ref = -np_log_softmax(np.array([2.0, 1.0]))[0]               # 只有第 1 行、前两个类别有效
assert np.isclose(float(loss_ok), ref, rtol=1e-6) and bool(jnp.all(jnp.isfinite(g_ok)))
assert np.allclose(probs_ok[1], 1 / 3)                        # 全 mask 行 → 均匀分布，无害
assert np.isnan(float(out["❌ 填 -inf"][0])) and bool(jnp.all(jnp.isnan(out["❌ 填 -inf"][1][1])))

def locate_nan(fn, *args, user_funcs=()):
    """开启 jax_debug_nans 运行 fn，返回 (异常信息, 用户代码所在行)。用 finally 保证关闭开关。"""
    import traceback
    jax.config.update("jax_debug_nans", True)
    jax.clear_caches()          # 0.12 实测：进程里用过 vmap 后，不清缓存会漏检（见下一个单元）
    try:
        jax.block_until_ready(fn(*args))
        return None, None
    except FloatingPointError as e:
        frames = [f for f in traceback.extract_tb(e.__traceback__) if f.name in user_funcs]
        where_ = f"{frames[-1].name}(): {frames[-1].line.strip()}" if frames else "（在 JAX 生成的反向计算里）"
        return str(e).splitlines()[0], where_
    finally:
        jax.config.update("jax_debug_nans", False)

def encoder(x):
    h = x * 2.0
    scale = jnp.log(h - 3.0)                 # h < 3 时 log(负数) = NaN
    return jnp.sum(scale * h)

def masked_log(x):
    return jnp.where(x > 0, jnp.log(x), 0.0)  # 前向的 NaN 被 where 丢掉，不会流到输出

x_bad = jnp.array([1.0, 5.0])
tests = [
    ("eager",                encoder,                       (x_bad,)),
    ("jit：输出含 NaN",       jax.jit(encoder),              (x_bad,)),
    ("grad(jit)：反向 NaN",   jax.grad(jax.jit(masked_log)), (jnp.array(0.0),)),
    ("jit：NaN 没到输出",     jax.jit(masked_log),           (jnp.array(-1.0),)),
    ("eager：同上函数",       masked_log,                    (jnp.array(-1.0),)),
    ("jit：只有 inf",         jax.jit(lambda x: 1.0 / x),    (jnp.array(0.0),)),
]
results = {}
for name, fn, args in tests:
    msg, loc = locate_nan(fn, *args, user_funcs={"encoder", "masked_log"})
    results[name] = msg
    print(f"[{name}]\n    {'未报错' if msg is None else 'FloatingPointError: ' + msg}" + (f"\n    定位到 → {loc}" if loc else ""))

assert results["eager"] and results["jit：输出含 NaN"] and results["grad(jit)：反向 NaN"]
assert results["jit：NaN 没到输出"] is None and results["eager：同上函数"] is not None
assert results["jit：只有 inf"] is None
assert jax.config.jax_debug_nans is False            # 开关已经恢复

# inf 要用 jax_debug_infs 单独抓
jax.config.update("jax_debug_infs", True)
jax.clear_caches()
try:
    jax.jit(lambda x: 1.0 / x)(jnp.array(0.0))
    caught_inf = None
except FloatingPointError as e:
    caught_inf = str(e).splitlines()[0]
finally:
    jax.config.update("jax_debug_infs", False)
print("jax_debug_infs:", caught_inf)
assert caught_inf is not None

from jax.experimental import checkify

checked = jax.jit(checkify.checkify(encoder, errors=checkify.float_checks))
err, out = checked(x_bad)
print("checkify 报告:", err.get())
err_ok, _ = checked(jnp.array([3.0, 5.0]))
print("正常输入:", err_ok.get())
err2, _ = checked(x_bad)                      # 命中缓存的第二次调用
print("再次传入坏输入:", err2.get())
assert err.get() is not None and err_ok.get() is None and err2.get() is not None
try:
    err.throw()
except Exception as e:
    print(f"err.throw() → {type(e).__name__}")