第 8 章　调试工具链

import os, sys, io, re, gc, time, json, logging, subprocess, textwrap, traceback, contextlib, functools, warnings
from functools import partial
from collections import Counter

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.monitoring as monitoring
from jax.experimental import checkify
import jax.test_util                      # 必须显式导入（见第 3 章），只 import jax 时访问 jax.test_util 会 AttributeError
from jax.test_util import check_grads

print("JAX", jax.__version__, "| 后端", jax.default_backend(), "| 设备", jax.devices())
CONFIG_AT_START = dict(jax.config.values)   # 记下开头的全部配置，最后检查调试开关是否都已复原


def to64(a):
    """转成 NumPy float64，用于对拍。"""
    return np.asarray(a, dtype=np.float64)


def show_error(fn):
    """运行 fn；报错时打印「类型: 第一行信息」并返回异常对象，不报错时返回 None。"""
    try:
        out = fn()
    except Exception as e:
        first = str(e).strip().splitlines()[0] if str(e).strip() else ""
        print(f"  → {type(e).__name__}: {first[:150]}")
        return e
    print("  → 没有报错，结果:", out)
    return None


def capture_stdout(fn):
    """运行 fn 并收集 stdout。jax.debug.print 是异步回调，所以最后要 effects_barrier() 等它打印完。"""
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        jax.block_until_ready(fn())
        jax.effects_barrier()
    return buf.getvalue()


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


class CompileCounter:
    """统计 with 块内的 XLA 后端编译次数；name="f" 时只数 jit(f)。"""
    EVENT = "/jax/core/compile/backend_compile_duration"

    def __init__(self, name=None):
        self.name, self.events = name, []
        self._listener = self._on_event          # 保存同一个绑定方法对象，注销时才找得到

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

    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 run_py(code, *args, env=None, stdin=None, timeout=300):
    """在全新的 Python 子进程里运行 code，返回 CompletedProcess。env 中的键会覆盖当前环境变量。"""
    full_env = {k: v for k, v in os.environ.items() if not k.startswith(("XLA_FLAGS", "JAX_"))}
    full_env["PYTHON_COLORS"] = "0"            # Python 3.14 的彩色 traceback 会在输出里混入 ANSI 转义码
    full_env.update(env or {})
    return subprocess.run([sys.executable, "-c", textwrap.dedent(code), *args], input=stdin,
                          capture_output=True, text=True, env=full_env, timeout=timeout)


def result_lines(proc):
    """取出子进程 stdout 中 'RESULT ' 之后的内容（它前面可能紧跟着调试器提示符 '(jdb) '），其他输出当作噪声忽略。"""
    return [l.split("RESULT ", 1)[1] for l in proc.stdout.splitlines() if "RESULT " in l]

@jax.jit
def stats(x):
    jax.debug.print("均值={m:.3f}  最大值={mx:.2f}  元素数={n}", m=x.mean(), mx=x.max(), n=x.size)
    jax.debug.print("位置参数：首={} 尾={}", x[0], x[-1])
    return x * 2

out = capture_stdout(lambda: stats(jnp.array([1.0, 2.0, 4.5])))
print(out)
assert "均值=2.500" in out and "最大值=4.50" in out and "元素数=3" in out
assert "位置参数：首=1.0 尾=4.5" in out

# ❌ 错误写法：f-string 在追踪时就被求值，打印出来的是 tracer，而且之后每次运行都打印这同一行
@jax.jit
def fstring_print(x):
    jax.debug.print(f"x = {x}")
    return x

try:
    out_bad = capture_stdout(lambda: fstring_print(jnp.float32(3.0)))
    print("实测输出:", out_bad.strip())
except Exception as e:                     # 如果 tracer 的字符串里恰好有花括号，还会直接报格式错误
    out_bad = f"{type(e).__name__}: {e}"
    print("实测报错:", out_bad)
assert "Tracer" in out_bad and "3.0" not in out_bad

# ✅ 正确写法：格式串是普通字符串，数值作为参数传进去
@jax.jit
def good_print(x):
    jax.debug.print("x = {x}", x=x)
    return x

out_good = capture_stdout(lambda: good_print(jnp.float32(3.0)))
print("实测输出:", out_good.strip())
assert out_good.strip() == "x = 3.0"

log_order = jax.debug.print(ordered=True)          # 第一段：只传选项，得到一个打印函数

@jax.jit
def stage_a(x):
    for i in range(3):                              # Python 循环在追踪时展开成 3 个打印
        log_order("A{i}", i=i)
    return x + 1

@jax.jit
def stage_b(x):
    log_order("B x={x}", x=x)
    return x * 2

out = capture_stdout(lambda: stage_b(stage_a(jnp.float32(1.0))))
print(out)
assert out.split() == ["A0", "A1", "A2", "B", "x=2.0"]

dbg_logger = logging.getLogger("jax._src.debugging")
old_level = dbg_logger.level
dbg_logger.setLevel(logging.INFO)              # 默认级别下 INFO 不会输出
try:
    with capture_jax_logs() as logs:
        jax.block_until_ready(jax.jit(lambda x: (jax.debug.log("loss = {x:.2f}", x=x), x * 2)[1])(jnp.float32(0.25)))
        jax.effects_barrier()
finally:
    dbg_logger.setLevel(old_level)
print("logging 收到:", logs.messages)
assert logs.messages == ["loss = 0.25"]

def body(x):
    jax.debug.print("看到 x = {x}", x=x)
    return jnp.sin(x) * x

def branchy(x):
    return lax.cond(x > 0,
                    lambda x: (jax.debug.print("正分支 {x}", x=x), x)[1],
                    lambda x: (jax.debug.print("负分支 {x}", x=x), -x)[1], x)

out_v = capture_stdout(lambda: jax.vmap(body)(jnp.array([1.0, 2.0, 3.0])))
out_g = capture_stdout(lambda: jax.grad(body)(2.0))
out_jg = capture_stdout(lambda: jax.jit(jax.grad(body))(2.0))
out_c = capture_stdout(lambda: jax.jit(branchy)(jnp.float32(1.0)))
out_vc = capture_stdout(lambda: jax.vmap(branchy)(jnp.array([1.0, -1.0])))
out_vo = capture_stdout(lambda: jax.vmap(lambda x: (jax.debug.print(ordered=True)("有序 {x}", x=x), x)[1])(jnp.arange(2.0)))
for name, o in [("vmap", out_v), ("grad", out_g), ("jit(grad)", out_jg), ("cond", out_c),
                ("vmap(cond)", out_vc), ("vmap + ordered", out_vo)]:
    print(f"{name:15s}", o.strip().replace("\n", " | "))

assert out_v.splitlines() == ["看到 x = 1.0", "看到 x = 2.0", "看到 x = 3.0"]
assert out_g.splitlines() == ["看到 x = 2.0"] and out_jg.splitlines() == ["看到 x = 2.0"]
assert out_c.splitlines() == ["正分支 1.0"]
assert sorted(out_vc.splitlines()) == sorted(["正分支 1.0", "正分支 -1.0", "负分支 1.0", "负分支 -1.0"])
assert out_vo.splitlines() == ["有序 0.0", "有序 1.0"]

GRAD_TAPS = []

@partial(jax.custom_vjp, nondiff_argnums=(1,))
def grad_tap(x, name):
    """恒等函数；反向传播经过这里时，把收到的梯度交给主机端记录下来。"""
    return x

def _grad_tap_fwd(x, name):
    return x, None

def _grad_tap_bwd(name, _, g):
    jax.debug.callback(lambda g: GRAD_TAPS.append((name, np.asarray(g))), g)
    return (g,)

grad_tap.defvjp(_grad_tap_fwd, _grad_tap_bwd)

def tap_loss(w, x):
    h = grad_tap(jnp.tanh(w * x), "h")
    return jnp.sum(h ** 2)

w_t, x_t = jnp.array([1.0, 2.0]), jnp.array([0.5, 0.1])
g_w = jax.jit(jax.grad(tap_loss))(w_t, x_t)
jax.effects_barrier()
name, g_h = GRAD_TAPS[-1]
print(f"{name} 收到的梯度:", g_h, "| 对 w 的梯度:", g_w)
h64 = np.tanh(to64(w_t) * to64(x_t))
assert name == "h" and np.allclose(g_h, 2 * h64, atol=1e-6)                       # ∂L/∂h = 2h
assert np.allclose(g_w, 2 * h64 * (1 - h64 ** 2) * to64(x_t), atol=1e-6)          # 探针不改变梯度

class ActivationLog:
    """用 jax.debug.callback 在 jit 内记录每层激活值的统计量。"""
    def __init__(self):
        self.records = []

    def record(self, name, x):
        def _host(m, s, frac_dead):              # 在主机端执行，参数是 NumPy 数组
            self.records.append({"layer": name, "mean": float(m), "std": float(s), "dead": float(frac_dead)})
        jax.debug.callback(_host, x.mean(), x.std(), jnp.mean(x <= 0))
        return x

act_log = ActivationLog()

def mlp_forward(params, x):
    for i, (w, b) in enumerate(params):
        x = x @ w + b
        if i < len(params) - 1:
            x = act_log.record(f"relu{i}", jax.nn.relu(x))
    return x

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
params_cb = [(jax.random.normal(k1, (8, 16)) / jnp.sqrt(8.0), jnp.zeros(16)),
             (jax.random.normal(k2, (16, 16)) / 4.0, jnp.zeros(16)),
             (jax.random.normal(k3, (16, 1)) / 4.0, jnp.zeros(1))]
x_cb = jax.random.normal(jax.random.key(1), (32, 8))
jax.block_until_ready(jax.jit(mlp_forward)(params_cb, x_cb))
jax.effects_barrier()
for r in act_log.records:
    print(f"  {r['layer']}: mean={r['mean']:.3f} std={r['std']:.3f} 死亡比例={r['dead']:.2f}")

# 对拍：eager 下用 NumPy 重新算第一层的统计量
h0 = np.maximum(to64(x_cb) @ to64(params_cb[0][0]), 0)
assert [r["layer"] for r in act_log.records] == ["relu0", "relu1"]
assert abs(act_log.records[0]["mean"] - h0.mean()) < 1e-5 and abs(act_log.records[0]["std"] - h0.std()) < 1e-5

# vmap 下：回调按样本逐个执行（4 个样本 × 2 层 = 8 条记录）
act_log.records.clear()
jax.block_until_ready(jax.vmap(mlp_forward, in_axes=(None, 0))(params_cb, x_cb[:4]))
jax.effects_barrier()
print("vmap 下的记录条数:", len(act_log.records))
assert len(act_log.records) == 8

BP_SCRIPT = """
import jax, jax.numpy as jnp

@jax.jit
def f(x):
    y = jnp.sin(x) * 2
    jax.debug.breakpoint()
    return y + 1

print("RESULT", f(jnp.arange(3.0)).tolist())
"""
proc = run_py(BP_SCRIPT, stdin="p y\np x\nc\n")
print(proc.stdout)
assert "Entering jdb" in proc.stdout and "(jdb)" in proc.stdout
assert "1.6829419" in proc.stdout                                   # p y：sin(1)*2，已经是具体数值
assert "Array([0., 1., 2.], dtype=float32)" in proc.stdout          # p x
res = json.loads(result_lines(proc)[0])                               # 输入 c 后程序继续执行完
assert np.allclose(res, np.sin(np.arange(3.0)) * 2 + 1, atol=1e-6)

COND_BP_SCRIPT = """
import jax, jax.numpy as jnp
from jax import lax

def stop_if_nan(x, y):            # 想看的变量作为参数传进来
    jax.debug.breakpoint()

@jax.jit
def f(x):
    y = jnp.log(x)
    lax.cond(jnp.any(jnp.isnan(y)), stop_if_nan, lambda x, y: None, x, y)
    return y

print("RESULT ok-run", f(jnp.array([1.0, 2.0])).tolist())     # 不停
print("RESULT nan-run", f(jnp.array([1.0, -2.0])).tolist())   # 停下来，查看后继续
"""
proc = run_py(COND_BP_SCRIPT, stdin="p y\np x\nc\n")
print(proc.stdout)
assert proc.stdout.count("Entering jdb") == 1                        # 只有含 NaN 的那次停下
assert "Array([ 0., nan], dtype=float32)" in proc.stdout and "Array([ 1., -2.], dtype=float32)" in proc.stdout
assert [l.split()[0] for l in result_lines(proc)] == ["ok-run", "nan-run"]

seen = []
with jax.disable_jit():
    @jax.jit
    def jitted_branch(x):
        seen.append(("jit 内参数类型", type(x).__name__))
        return x if x > 0 else -x                       # Python if：只有 disable_jit 下才合法
    r1 = jitted_branch(jnp.array(-3.0))
    lax.scan(lambda c, x: (seen.append(("scan 体内 carry", type(c).__name__)) or (c + x, c)), jnp.float32(0), jnp.arange(2.0))
    lax.fori_loop(0, 2, lambda i, v: (seen.append(("fori 下标", type(i).__name__)) or v + i), 0)
    jax.vmap(lambda x: seen.append(("vmap 内参数", type(x).__name__)) or x)(jnp.arange(2.0))
for s in seen:
    print("  ", *s)
kinds = dict(seen)
assert float(r1) == 3.0 and kinds["jit 内参数类型"] == "ArrayImpl"
assert kinds["fori 下标"] == "int" and kinds["vmap 内参数"] == "BatchTracer"

def bisect_numeric(fn, *args, ref):
    """同一个函数分别在 jit / disable_jit 下运行，和参考值比较最大误差。"""
    err_jit = float(np.abs(to64(jax.jit(fn)(*args)) - ref).max())
    with jax.disable_jit():
        err_eager = float(np.abs(to64(fn(*args)) - ref).max())
    print(f"    jit 误差 = {err_jit:.3g} | disable_jit 误差 = {err_eager:.3g}")
    return err_jit, err_eager

W_bug = jax.random.normal(jax.random.key(1), (128, 128)) / jnp.sqrt(128.0)
X_bug = jax.random.normal(jax.random.key(2), (50, 128))
ref_bug = (to64(W_bug) @ to64(X_bug).T).sum(0)                     # NumPy float64 参考值

# ❌ 错误写法：权重被 jit 闭包捕获为常量（jaxlib 0.12 在 macOS / Apple Silicon 的 CPU 后端上会悄悄算错）
def closure_version(X):
    return jax.vmap(lambda x: jnp.sum(W_bug @ x))(X)

try:
    err_jit, err_eager = bisect_numeric(closure_version, X_bug, ref=ref_bug)
except Exception as e:
    print("实测报错:", type(e).__name__, e)
    raise
assert err_eager < 1e-3
import platform, sys
if sys.platform == "darwin" and platform.machine() == "arm64":   # 在 Apple Silicon 上必须复现；不复现说明可能已修复
    assert err_jit > 1.0, "macOS arm64 上没有复现：可能已在新版 jaxlib 修复，请更新相关说明"
if err_jit > 1.0:
    print("jit 错、eager 对 → 问题出在编译这一层")
else:
    print("本机没有触发这个 bug，jit 的结果也是对的（它在 macOS / Apple Silicon 上能复现，在 Linux x86 上测试时没有出现）。二分的思路不变。")

# ✅ 正确写法：权重作为函数参数传入
def param_version(W, X):
    return jax.vmap(lambda x: jnp.sum(W @ x))(X)

err_jit, err_eager = bisect_numeric(param_version, W_bug, X_bug, ref=ref_bug)
assert err_jit < 1e-3 and err_eager < 1e-3

SCALE = 2.0                                       # 模块级「配置」

def scaled(x):
    return x * SCALE

jaxpr_txt = str(jax.make_jaxpr(scaled)(1.0))
print(jaxpr_txt)
assert "mul a 2.0:f32[]" in jaxpr_txt             # SCALE 以字面量形式出现：它被烘焙进去了

# ❌ 错误写法：配置放在全局变量里，jit 之后再改它
scaled_jit = jax.jit(scaled)
before = float(scaled_jit(1.0))
SCALE = 10.0
try:
    after = float(scaled_jit(1.0))
    print(f"修改 SCALE 前 = {before}，修改后 = {after}（没有报错，但用的仍是旧值）")
except Exception as e:
    print("实测报错:", type(e).__name__, e)
    raise
assert before == after == 2.0

# ✅ 正确写法：配置作为参数传入；jaxpr 里它变成了输入变量 b，而不是字面量
def scaled_ok(x, scale):
    return x * scale

print(jax.make_jaxpr(scaled_ok)(1.0, SCALE))
scaled_ok_jit = jax.jit(scaled_ok)
assert float(scaled_ok_jit(1.0, 2.0)) == 2.0 and float(scaled_ok_jit(1.0, 10.0)) == 10.0
SCALE = 2.0

def scoped_mlp(params, x):
    with jax.named_scope("encoder"):
        h = jax.nn.relu(x @ params["w1"])
    with jax.named_scope("head"):
        out = h @ params["w2"]
        loss = jnp.mean(out ** 2)
    return loss

p_ns = {"w1": jnp.ones((4, 8)) * 0.1, "w2": jnp.ones((8, 2)) * 0.1}
x_ns = jnp.ones((3, 4))
jf_ns = jax.jit(scoped_mlp)

# ① jaxpr：每个等式的名字栈
traced = jf_ns.trace(p_ns, x_ns)
stacks = [str(e.source_info.name_stack) for e in traced.jaxpr.jaxpr.eqns]
print("① jaxpr 名字栈:", stacks)

# ② StableHLO：默认文本里没有名字，debug_info=True 才有
lowered = traced.lower()
plain, dbg = lowered.as_text(), lowered.as_text(debug_info=True)
print("② StableHLO 默认文本含 'encoder'？", "encoder" in plain, "| debug_info=True 时：",
      [l.strip() for l in dbg.splitlines() if "encoder/" in l][:2])

# ③ 优化后 HLO：metadata 里的 op_name
compiled = lowered.compile()
def hlo_ops_in_scope(hlo_text, scope):
    """从优化后 HLO 里找出 op_name 含 /scope/ 的指令，返回 (指令名, op_name)。"""
    pat = re.compile(r"%([\w.\-]+) = .*?op_name=\"([^\"]*/" + re.escape(scope) + r"/[^\"]*)\"")
    return [m.groups() for m in pat.finditer(hlo_text)]
enc_ops = hlo_ops_in_scope(compiled.as_text(), "encoder")
print("③ 优化后 HLO 中属于 encoder 的指令:")
for name, op in enc_ops:
    print(f"     %{name:20s} ← {op}")

assert stacks.count("encoder") >= 2 and stacks.count("head") >= 3
assert "encoder" not in plain and "jit(scoped_mlp)/encoder/dot_general" in dbg
assert any(op.endswith("encoder/dot_general") for _, op in enc_ops)
assert np.isclose(float(compiled(p_ns, x_ns)), float(np.mean((np.maximum(to64(x_ns) @ to64(p_ns["w1"]), 0) @ to64(p_ns["w2"])) ** 2)))

print("类型:", type(traced).__name__, "→", type(lowered).__name__, "→", type(compiled).__name__)
print("Traced.out_info:", traced.out_info, "| Compiled.input_shardings[0] 的键:", sorted(compiled.input_shardings[0][0]))
assert isinstance(traced, jax.stages.Traced) and isinstance(lowered, jax.stages.Lowered)
assert isinstance(compiled, jax.stages.Compiled)
assert traced.out_info.shape == () and traced.out_info.dtype == jnp.float32

with jax.debug_nans(True):                          # with 写法不受下面的 vmap 漏检问题影响，不用清缓存
    print("块内:", jax.config.jax_debug_nans)
    e_nan = show_error(lambda: jax.jit(lambda x: jnp.log(x))(jnp.array(-1.0)))
with jax.debug_infs(True):
    e_inf = show_error(lambda: jax.jit(lambda x: 1.0 / x)(jnp.array(0.0)))
print("块外:", jax.config.jax_debug_nans, jax.config.jax_debug_infs)
assert isinstance(e_nan, FloatingPointError) and "nan" in str(e_nan)
assert isinstance(e_inf, FloatingPointError) and "inf" in str(e_inf)
assert jax.config.jax_debug_nans is False and jax.config.jax_debug_infs is False

NAN_PROBE = """
import sys, jax, jax.numpy as jnp
mode = sys.argv[1]
if mode == "with":
    jax.debug_nans(True).__enter__()               # 等价于把下面整段放进 with 块
else:
    jax.config.update("jax_debug_nans", True)
f_old = jax.jit(lambda x: jnp.log(x))
f_old(jnp.array(2.0)).block_until_ready()          # 先正常执行一次
jnp.log(jnp.array(2.0))                            # eager 的 log 也先执行一次
if mode in ("vmap", "clear", "with"):
    jax.vmap(lambda x: x * 2)(jnp.ones(3))         # 一次无关的 vmap
if mode == "clear":
    jax.clear_caches()                             # 规避方法
cases = {
    "已执行过的 jit 函数": lambda: f_old(jnp.array(-1.0)),
    "已执行过的 eager 原语 log": lambda: jnp.log(jnp.array(-1.0)),
    "没执行过的 eager 原语 sqrt": lambda: jnp.sqrt(jnp.array(-1.0)),
    "新建的 jit 函数": lambda: jax.jit(lambda x: jnp.arccos(x))(jnp.array(3.0)),
}
for name, fn in cases.items():
    try:
        fn(); print("RESULT", name, "漏检")
    except FloatingPointError:
        print("RESULT", name, "捕获")
"""
nan_table = {}
for mode in ["none", "vmap", "clear", "with"]:
    for line in result_lines(run_py(NAN_PROBE, mode)):
        name, verdict = line.rsplit(" ", 1)
        nan_table.setdefault(name, {})[mode] = verdict
print(f"{'场景':24s} {'没用过 vmap':10s} {'用过 vmap':10s} {'vmap+clear':10s} {'with 写法+vmap':10s}")
for name, row in nan_table.items():
    print(f"{name:22s} {row['none']:12s} {row['vmap']:12s} {row['clear']:12s} {row['with']}")

assert all(row["none"] == row["clear"] == row["with"] == "捕获" for row in nan_table.values())
assert nan_table["已执行过的 jit 函数"]["vmap"] == "漏检" and nan_table["已执行过的 eager 原语 log"]["vmap"] == "漏检"
assert nan_table["没执行过的 eager 原语 sqrt"]["vmap"] == "捕获" and nan_table["新建的 jit 函数"]["vmap"] == "捕获"

TRIGGER_PROBE = """
import sys, jax, jax.numpy as jnp
from jax.experimental import checkify
jax.config.update("jax_debug_nans", True)
f = jax.jit(lambda x: jnp.log(x))
f(jnp.array(2.0)).block_until_ready()
m = sys.argv[1]
if m == "grad": jax.grad(lambda x: x * 2)(1.0)
if m == "scan": jax.lax.scan(lambda c, x: (c + x, x), 0.0, jnp.ones(3))
if m == "make_jaxpr": jax.make_jaxpr(lambda x: x * 2)(1.0)
if m == "checkify": checkify.checkify(lambda x: x * 2)(1.0)
if m == "config.update": jax.config.update("jax_debug_nans", True)
if m == "make_jaxpr(vmap)": jax.make_jaxpr(jax.vmap(lambda x: x * 2))(jnp.ones(3))
if m == "jit(vmap)": jax.jit(jax.vmap(lambda x: x * 2))(jnp.ones(3))
try:
    f(jnp.array(-1.0)); print("RESULT 漏检")
except FloatingPointError:
    print("RESULT 捕获")
"""
trig = {m: result_lines(run_py(TRIGGER_PROBE, m))[0]
        for m in ["grad", "scan", "make_jaxpr", "checkify", "config.update", "make_jaxpr(vmap)", "jit(vmap)"]}
for m, v in trig.items():
    print(f"  先做 {m:18s} → 旧 jit 函数的 NaN：{v}")
assert [m for m, v in trig.items() if v == "漏检"] == ["make_jaxpr(vmap)", "jit(vmap)"]

def reg_loss(p, x):
    return jnp.sum((x @ p["w"] + p["b"]) ** 2) + 1e-3 * jnp.linalg.norm(p["b"])   # b 初始化为 0

p_reg = {"w": jnp.ones((3, 2)), "b": jnp.zeros(2)}
x_reg = jnp.ones((4, 3))
with jax.debug_nans(True):
    jax.clear_caches()
    err = show_error(lambda: jax.grad(reg_loss)(p_reg, x_reg))
user_frames = [f for f in traceback.extract_tb(err.__traceback__) if "site-packages" not in f.filename]
print("traceback 里的用户代码帧:", [(f.name, f.line) for f in user_frames])

# 逐叶子检查：哪个参数的梯度含 NaN
grads = jax.grad(reg_loss)(p_reg, x_reg)
bad_leaves = [jax.tree_util.keystr(path) for path, g in jax.tree_util.tree_leaves_with_path(grads)
              if not bool(jnp.all(jnp.isfinite(g)))]
print("梯度含 NaN 的叶子:", bad_leaves)
assert isinstance(err, FloatingPointError) and "mul" in str(err)
assert "reg_loss" not in [f.name for f in user_frames]        # 没有一帧落在 reg_loss 里
assert bad_leaves == ["['b']"]

def risky(x, i):
    checkify.check(jnp.all(x > 0), "x 必须为正，最小值 = {m}", m=jnp.min(x))
    y = jnp.log(x)                  # x ≤ 0 → NaN
    z = x[i]                        # i 越界 → JAX 默认夹紧下标，不报错
    w = 7 // i                      # i == 0 → 整数除零
    return y.sum() + z + w

x_ok = jnp.array([1.0, 2.0, 3.0])
inputs = {"正常": (x_ok, 1), "x 含负数": (jnp.array([-1.0, 2.0, 3.0]), 1), "下标越界": (x_ok, 5), "除零": (x_ok, 0)}
sets = {"user": checkify.user_checks, "nan": checkify.nan_checks, "index": checkify.index_checks,
        "div": checkify.div_checks, "all": checkify.all_checks}
table = {}
for sname, errs in sets.items():
    checked = jax.jit(checkify.checkify(risky, errors=errs))
    for iname, args in inputs.items():
        err, _ = checked(*args)
        table[(sname, iname)] = err.get()
print(f"{'检查集合':8s}" + "".join(f"{n:>10s}" for n in inputs))
for sname in sets:
    print(f"{sname:10s}" + "".join(f"{'✔ 抓到' if table[(sname, n)] else '—':>10s}" for n in inputs))
print("\n示例消息：")
for key in [("user", "x 含负数"), ("nan", "x 含负数"), ("index", "下标越界"), ("div", "除零")]:
    print(f"  {key}: {table[key].strip()[:90]}")

assert all(table[(s, "正常")] is None for s in sets)
assert "x 必须为正，最小值 = -1.0" in table[("user", "x 含负数")]
assert "nan generated by primitive: log" in table[("nan", "x 含负数")]
assert "out-of-bounds" in table[("index", "下标越界")] and table[("user", "下标越界")] is None
assert "division by zero" in table[("div", "除零")]
assert all(table[("all", n)] is not None for n in inputs if n != "正常")
clamped = float(jax.jit(lambda x, i: x[i])(x_ok, 5))
print("\n不开检查时 x[5] 的实际结果:", clamped, "（下标被夹紧成 2，静默错误）")
assert clamped == 3.0

# ❌ 错误写法：在 jit 里直接用 checkify.check，却没有用 checkify.checkify 包装
def check_only(x):
    checkify.check(x > 0, "x 必须为正")
    return jnp.log(x)

try:
    jax.jit(check_only)(1.0)
    raised = None
except Exception as e:
    raised = e
    print(f"实测报错: {type(e).__name__}: {str(e).splitlines()[0][:120]}")
assert isinstance(raised, ValueError) and "not functionalized" in str(raised)

# ✅ 正确写法：先 checkify 再 jit；拿到 err 后，用 err.get() 查看或 err.throw() 抛出
checked_only = jax.jit(checkify.checkify(check_only))
err_ok, _ = checked_only(1.0)
err_bad, _ = checked_only(-1.0)
print("正常输入:", err_ok.get(), "| 坏输入:", err_bad.get())
try:
    err_bad.throw()
except Exception as e:
    thrown = e
    print(f"err.throw() → {type(e).__name__}（继承链 {[c.__name__ for c in type(e).__mro__[:3]]}）")
assert err_ok.get() is None and "x 必须为正" in err_bad.get()
assert isinstance(thrown, checkify.JaxRuntimeError) and isinstance(thrown, ValueError)

def pos_log(x):
    checkify.check(x > 0, "x 必须为正，实际为 {v}", v=x)
    return jnp.log(x)

xs_mixed = jnp.array([1.0, -2.0, 3.0, -4.0])
r = {}
r["jit(checkify)"] = jax.jit(checkify.checkify(pos_log))(jnp.float32(-1.0))[0].get()
r["checkify(jit)"] = checkify.checkify(jax.jit(pos_log))(jnp.float32(-1.0))[0].get()
r["vmap(checkify)"] = jax.vmap(checkify.checkify(pos_log))(xs_mixed)[0].get()
r["checkify(vmap)"] = checkify.checkify(jax.vmap(pos_log))(xs_mixed)[0].get()
bad_where = lambda x: jnp.sum(jnp.where(x > 0, jnp.sqrt(x), 0.0))          # 单 where：反向会出 NaN
r["checkify(grad)"] = checkify.checkify(jax.grad(bad_where), errors=checkify.nan_checks)(jnp.array([0.0, 1.0]))[0].get()
g_inside = jax.grad(lambda x: checkify.checkify(pos_log)(x)[1])(jnp.float32(-1.0))
def scan_body(c, x):
    checkify.check(x < 10, "第 {i} 步的值太大：{x}", i=c[1], x=x)
    return (c[0] + x, c[1] + 1), None
r["scan"] = jax.jit(checkify.checkify(lambda xs: lax.scan(scan_body, (0.0, 0), xs)[0]))(jnp.array([1.0, 20.0, 30.0]))[0].get()
r["float_checks 与 inf"] = checkify.checkify(jnp.exp, errors=checkify.float_checks)(jnp.float32(1000.0))[0].get()
@jax.jit
def dbg_fn(x):
    checkify.debug_check(x > 0, "debug_check 失败：{x}", x=x)
    return x * 2
dbg_plain = float(dbg_fn(-1.0))
r["debug_check + checkify"] = checkify.checkify(dbg_fn)(-1.0)[0].get()
for k, v in r.items():
    print(f"  {k:22s} → {v.strip() if v else None}".replace("\n", " ⏎ "))
print("  grad(checkify(f)[1]) 的梯度:", g_inside, "| debug_check 未包装时的结果:", dbg_plain)

assert r["jit(checkify)"] == r["checkify(jit)"] and "-1.0" in r["jit(checkify)"]
assert "at mapped index 1" in r["vmap(checkify)"] and "at mapped index 3" in r["vmap(checkify)"]
assert "mapped index" not in r["checkify(vmap)"] and r["checkify(vmap)"].count("必须为正") == 1
assert "nan generated by primitive: mul" in r["checkify(grad)"]
assert float(g_inside) == -1.0                                  # 能求导，但检查被丢掉了
assert "第 1 步的值太大：20.0" in r["scan"]
assert r["float_checks 与 inf"] is None
assert dbg_plain == -2.0 and "debug_check 失败" in r["debug_check + checkify"]

e_py = show_error(lambda: checkify.checkify(pos_log)(-1.0))
e_arr, _ = checkify.checkify(pos_log)(jnp.asarray(-1.0))
assert isinstance(e_py, TypeError) and "PyTrees of arrays" in str(e_py)
assert "实际为 -1.0" in e_arr.get()

# ❌ 错误写法：custom_vjp 的反向规则写错了符号（softplus 的导数应为 sigmoid(x)）
@jax.custom_vjp
def softplus_bad(x):
    return jnp.logaddexp(x, 0.0)
softplus_bad.defvjp(lambda x: (softplus_bad(x), x), lambda x, g: (g * jax.nn.sigmoid(-x),))

x_cg = jnp.array([-1.0, 0.5, 2.0])
try:
    check_grads(softplus_bad, (x_cg,), order=1, modes=["rev"])
    cg_err = None
except AssertionError as e:
    cg_err = e
    print("check_grads 发现错误:", " | ".join(str(e).strip().splitlines()[:2]))
e_fwd = show_error(lambda: check_grads(softplus_bad, (x_cg,), order=1, modes=["fwd"]))
assert cg_err is not None and "VJP" in str(cg_err)
assert isinstance(e_fwd, TypeError) and "custom_vjp" in str(e_fwd)

# ✅ 正确写法：反向规则用 sigmoid(x)，check_grads 安静通过（通过时不输出任何东西）
@jax.custom_vjp
def softplus_ok(x):
    return jnp.logaddexp(x, 0.0)
softplus_ok.defvjp(lambda x: (softplus_ok(x), x), lambda x, g: (g * jax.nn.sigmoid(x),))

check_grads(softplus_ok, (x_cg,), order=1, modes=["rev"])
assert np.allclose(jax.grad(lambda x: softplus_ok(x).sum())(x_cg), 1 / (1 + np.exp(-to64(x_cg))), atol=1e-6)
print("check_grads 通过 ✔")

E = jax.errors
for cls in [E.ConcretizationTypeError, E.TracerBoolConversionError, E.TracerIntegerConversionError,
            E.TracerArrayConversionError, E.NonConcreteBooleanIndexError, E.UnexpectedTracerError]:
    print(f"  {cls.__name__:30s} → {[c.__name__ for c in cls.__mro__[1:4]]}")
assert issubclass(E.TracerBoolConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerIntegerConversionError, E.ConcretizationTypeError)
assert not issubclass(E.TracerArrayConversionError, E.ConcretizationTypeError)
assert all(issubclass(c, E.JAXTypeError) and issubclass(c, TypeError) for c in
           [E.ConcretizationTypeError, E.TracerIntegerConversionError, E.TracerArrayConversionError, E.UnexpectedTracerError])
assert issubclass(E.NonConcreteBooleanIndexError, IndexError)

def clip_positive(x, limit):
    if x > limit:                  # 依赖 x 和 limit 的数值
        return limit
    return x

e = show_error(lambda: jax.jit(clip_positive)(jnp.float32(3.0), 1.0))
msg = " ".join(str(e).split())                       # 原信息有换行，先压成一行再匹配
m = re.search(r"The error occurred while tracing the function (\S+) at \S+ for jit\..*?"
              r"depends on the values? of the arguments? (.+?)\.", msg)
print("出问题的函数:", m.group(1), "| 罪魁祸首参数:", m.group(2))
assert isinstance(e, E.TracerBoolConversionError)
assert m.group(1) == "clip_positive" and m.group(2) == "x and limit"   # 两个参数都参与了比较

# ❌ 错误写法：Python if 依赖追踪值
@jax.jit
def relu_if(x):
    if x > 0:
        return x
    return 0.0 * x

try:
    relu_if(jnp.float32(2.0))
    caught = None
except Exception as e:
    caught = e
    print(f"实测报错: {type(e).__name__}: {str(e).splitlines()[0]}")
assert isinstance(caught, E.TracerBoolConversionError)

# ✅ 正确写法：逐元素用 jnp.where；需要只执行一个分支时用 lax.cond
@jax.jit
def relu_where(x):
    return jnp.where(x > 0, x, 0.0)

@jax.jit
def relu_cond(x):
    return lax.cond(x > 0, lambda x: x, lambda x: 0.0 * x, x)

xs = jnp.array([-2.0, 0.0, 3.0])
assert np.array_equal(relu_where(xs), [0.0, 0.0, 3.0])
assert [float(relu_cond(v)) for v in xs] == [0.0, 0.0, 3.0]
print("jnp.where:", relu_where(xs), "| lax.cond:", [float(relu_cond(v)) for v in xs])

# ❌ 错误写法：用追踪值做 range 的次数、做 Python 列表的下标
def repeat_add(x, n):
    for _ in range(n):
        x = x + 1.0
    return x

def pick(i):
    return [10.0, 20.0, 30.0][i]

for name, fn in [("range(n)", lambda: jax.jit(repeat_add)(1.0, 3)), ("list[i]", lambda: jax.jit(pick)(1))]:
    try:
        fn()
        caught = None
    except Exception as e:
        caught = e
        print(f"{name:9s} 实测报错: {type(e).__name__}: {str(e).splitlines()[0]}")
    assert isinstance(caught, E.TracerIntegerConversionError)

# ✅ 正确写法：次数决定计算结构 → 静态参数；或用 fori_loop；列表先转成数组再索引
repeat_static = jax.jit(repeat_add, static_argnames="n")
repeat_fori = jax.jit(lambda x, n: lax.fori_loop(0, n, lambda _, v: v + 1.0, x))
pick_arr = jax.jit(lambda i: jnp.asarray([10.0, 20.0, 30.0])[i])
assert float(repeat_static(1.0, n=3)) == 4.0 and float(repeat_fori(1.0, 3)) == 4.0 and float(pick_arr(1)) == 20.0
print("static:", repeat_static(1.0, n=3), "| fori_loop:", repeat_fori(1.0, 3), "| 数组索引:", pick_arr(1))