第 9 章　导出与部署

import os, re, sys, json, time, shutil, textwrap, subprocess, tempfile, importlib.util
from importlib.metadata import requires, version
from pathlib import Path

import numpy as np
import jax
import jax.numpy as jnp
from jax import export
import jaxlib, flax, optax
from flax import nnx
import orbax.checkpoint as ocp
import absl.logging
absl.logging.set_verbosity(absl.logging.ERROR)        # 只保留 Orbax 的错误日志，避免 INFO/WARNING 刷屏

print("Python", sys.version.split()[0], "| JAX", jax.__version__, "| jaxlib", jaxlib.__version__,
      "| Flax", flax.__version__, "| Optax", optax.__version__, "| Orbax", version("orbax-checkpoint"))
print("设备:", jax.devices())

# 导出产物目录：与本 notebook 同级（notebooks/进阶/export_artifacts/），每次运行先清空
ART = Path("export_artifacts").resolve()
if ART.exists():
    shutil.rmtree(ART)
ART.mkdir(parents=True)
SCRATCH = Path(tempfile.mkdtemp(prefix="ext03_"))     # 中间文件放临时目录，不进仓库
print("产物目录:", ART.name + "/")

def run_py(src, *args, env_extra=None, cwd=None):
    """在全新子进程（python -I 隔离模式）里运行脚本 src，返回 stdout；失败时抛出带 stderr 的异常。"""
    env = {**os.environ, **(env_extra or {})}
    r = subprocess.run([sys.executable, "-I", "-c", textwrap.dedent(src), *map(str, args)],
                       capture_output=True, text=True, env=env, cwd=str(cwd or SCRATCH))
    if r.returncode != 0:
        raise RuntimeError(f"子进程失败（返回码 {r.returncode}）：\n{r.stderr[-3000:]}")
    return r.stdout

def sds(tree):
    """数组 PyTree → 同结构的 ShapeDtypeStruct（只保留形状与 dtype）。"""
    return jax.tree.map(lambda a: jax.ShapeDtypeStruct(np.shape(a), a.dtype), tree)

def max_diff(a, b):
    """两个同结构 PyTree 的最大逐元素绝对误差（float64 比较）。"""
    la, lb = jax.tree.leaves(a), jax.tree.leaves(b)
    assert len(la) == len(lb)
    return max(float(np.max(np.abs(np.asarray(x, np.float64) - np.asarray(y, np.float64)))) for x, y in zip(la, lb))

def tree_bytes(tree):
    return sum(np.asarray(x).nbytes for x in jax.tree.leaves(tree))

def dir_bytes(p):
    p = Path(p)
    return p.stat().st_size if p.is_file() else sum(f.stat().st_size for f in p.rglob("*") if f.is_file())

def main_signature(exp):
    """取出导出模块 @main 的签名行（看参数与返回值类型）。"""
    return next(l.strip() for l in exp.mlir_module().splitlines() if "func.func public @main" in l)

# ❌ 错误写法：环境里没装 flatbuffers 就调用 serialize()——jax 不会自动安装它
# 不真的卸载包：在子进程里令 sys.modules["flatbuffers"] = None，此后 import flatbuffers 必然抛 ImportError，以此模拟"未安装"
try:
    out = run_py("""
        import sys
        sys.modules["flatbuffers"] = None
        import jax, jax.numpy as jnp
        from jax import export
        exp = export.export(jax.jit(jnp.sin))(jnp.ones(3))
        print("导出本身成功，platforms =", exp.platforms)
        try:
            exp.serialize()
            print("RESULT 序列化成功")
        except ImportError as e:
            print("RESULT ImportError:", e)
    """)
    print("实测：\n" + out)
except Exception as e:
    print("实测报错：", type(e).__name__, e)
assert "导出本身成功" in out
assert "RESULT ImportError: Please install 'flatbuffers' in order to use Exported serialization" in out

# ✅ 正确写法：单独安装 flatbuffers（pip install flatbuffers），并写进 requirements.txt
import flatbuffers
deps = (requires("jax") or []) + (requires("jaxlib") or [])
print("flatbuffers 版本:", flatbuffers.__version__)
print("jax/jaxlib 声明的依赖中含 flatbuffers 吗？", any(d.lower().startswith("flatbuffers") for d in deps))
assert not any(d.lower().startswith("flatbuffers") for d in deps)      # 确实不是 jax 的依赖

blob = export.export(jax.jit(jnp.sin))(jnp.ones(3)).serialize()
print("serialize() 正常：", type(blob).__name__, len(blob), "字节")
assert isinstance(blob, bytearray) and len(blob) > 0

def linear_apply(params, x):
    """最小的"模型"：程序只描述计算，权重从参数传入。"""
    return x @ params["w"] + params["b"]

k1, k2, k3 = jax.random.split(jax.random.key(0), 3)
params_a = {"w": jax.random.normal(k1, (4, 2)), "b": jnp.zeros(2)}          # 第一份权重
params_b = {"w": jax.random.normal(k2, (4, 2)), "b": jnp.ones(2)}           # 第二份权重（比如重新训练过）
x_demo = jax.random.normal(k3, (3, 4))

# 程序：只导出一次（只用到形状和 dtype）
prog = export.export(jax.jit(linear_apply))(sds(params_a), sds(x_demo))
# 数据：两份权重分别用 Orbax 保存
ckptr = ocp.StandardCheckpointer()
for name, p in [("a", params_a), ("b", params_b)]:
    ckptr.save(SCRATCH / f"linear_{name}", p)
ckptr.wait_until_finished()

for name, p in [("a", params_a), ("b", params_b)]:
    restored = ckptr.restore(SCRATCH / f"linear_{name}", sds(p))           # 用抽象模板恢复
    y = prog.call(restored, x_demo)
    print(f"权重 {name}: 导出程序的输出 = {np.round(np.asarray(y[0]), 4)}")
    assert max_diff(y, linear_apply(p, x_demo)) < 1e-6
assert max_diff(prog.call(params_a, x_demo), prog.call(params_b, x_demo)) > 0.1
print("同一个导出程序 + 两份 Orbax 权重 → 两个不同模型，都与原函数一致 ✓")

def mlp_apply(params, x):
    """两层 MLP 的纯函数版本。权重 params 是参数，不是闭包常量。"""
    h = jnp.tanh(x @ params["w1"] + params["b1"])
    return h @ params["w2"] + params["b2"]

def init_mlp(key, din=4, dh=8, dout=3):
    k1, k2 = jax.random.split(key)
    return {"w1": jax.random.normal(k1, (din, dh)) / np.sqrt(din), "b1": jnp.zeros(dh),
            "w2": jax.random.normal(k2, (dh, dout)) / np.sqrt(dh), "b2": jnp.zeros(dout)}

params = init_mlp(jax.random.key(0))
x2 = jax.random.normal(jax.random.key(1), (2, 4))

exp = export.export(jax.jit(mlp_apply))(params, x2)
print("fun_name     :", exp.fun_name)
print("in_tree      :", exp.in_tree)
print("in_avals     :", exp.in_avals)
print("out_avals    :", exp.out_avals)
print("platforms    :", exp.platforms, "| nr_devices:", exp.nr_devices, "| calling_convention_version:", exp.calling_convention_version)
print("@main 签名   :", main_signature(exp)[:200], "...")

assert exp.fun_name == "mlp_apply" and exp.platforms == ("cpu",) and exp.nr_devices == 1
assert [a.shape for a in exp.in_avals] == [(8,), (3,), (4, 8), (8, 3), (2, 4)]       # b1, b2, w1, w2, x
assert exp.out_avals[0].shape == (2, 3) and exp.in_tree.num_leaves == 5
assert "stablehlo.dot_general" in exp.mlir_module() and "stablehlo.tanh" in exp.mlir_module()
assert max_diff(exp.call(params, x2), mlp_apply(params, x2)) < 1e-6

# 只给 ShapeDtypeStruct 也能导出，得到完全相同的程序
exp_abs = export.export(jax.jit(mlp_apply))(sds(params), jax.ShapeDtypeStruct((2, 4), jnp.float32))
assert exp_abs.mlir_module_serialized == exp.mlir_module_serialized
print("只用 ShapeDtypeStruct 导出：字节码与用真实数组导出的完全相同 ✓")

# ❌ 错误写法：直接导出普通 Python 函数（没有用 jax.jit 包装）
try:
    export.export(mlp_apply)(params, x2)
    print("实测：可以运行")
except ValueError as e:
    print("实测报错：ValueError:", e)
    err_nojit = str(e)
assert "must be the result of `jit`" in err_nojit

# ✅ 正确写法：export.export 只接受 jax.jit 的结果；jit 的 static_argnums、donate 等选项也一并生效
exp_ok = export.export(jax.jit(mlp_apply))(params, x2)
assert max_diff(exp_ok.call(params, x2), mlp_apply(params, x2)) < 1e-6
print("export.export(jax.jit(f)) ✓")

blob = exp.serialize()
exp_path = SCRATCH / "mlp_static.jaxexport"
exp_path.write_bytes(blob)
io_path = SCRATCH / "mlp_static_io.npz"
np.savez(io_path, x=np.asarray(x2), **{k: np.asarray(v) for k, v in params.items()})
print(f"导出文件 {exp_path.name}: {len(blob)} 字节；文件头 {bytes(blob[:8])!r}")

sub_src = """
    import sys, numpy as np, jax
    from jax import export
    exp = export.deserialize(open(sys.argv[1], "rb").read())     # 只读字节，不 import 任何模型代码
    d = np.load(sys.argv[2])
    params = {k: d[k] for k in ("w1", "b1", "w2", "b2")}
    np.save(sys.argv[3], np.asarray(exp.call(params, d["x"])))
    print("子进程 fun_name =", exp.fun_name, "| in_avals =", exp.in_avals)
    print("子进程 sys.path[0] =", sys.path[0])
"""
out_path = SCRATCH / "mlp_static_out.npy"
print(run_py(sub_src, exp_path, io_path, out_path))

y_sub = np.load(out_path)
print("子进程输出与原函数的最大差:", max_diff(y_sub, mlp_apply(params, x2)))
assert max_diff(y_sub, mlp_apply(params, x2)) < 1e-6
for token in ("mlp_apply", "tanh", "@", "def "):                 # 子进程脚本里没有任何计算逻辑
    assert token not in sub_src, token

# 同一进程里 deserialize 回来的对象与原对象签名一致
exp_back = export.deserialize(blob)
assert exp_back.in_avals == exp.in_avals and exp_back.out_avals == exp.out_avals and exp_back.in_tree == exp.in_tree
print("序列化 → 文件 → 新子进程 deserialize → call，结果与原函数一致 ✓")

# ❌ 错误写法：把权重作为闭包常量烘焙进导出程序（程序与数据绑死、文件膨胀，还会触发 CPU 后端 bug）
rng = np.random.default_rng(0)
n = 128
W_np = rng.standard_normal((n, n)).astype(np.float32)
X_np = rng.standard_normal((n, n)).astype(np.float32)
ref = float(np.sum(W_np.astype(np.float64) @ X_np.astype(np.float64)))      # float64 参考值
W_const = jnp.asarray(W_np)
try:
    exp_closure = export.export(jax.jit(lambda x: jnp.sum(W_const @ x)))(jax.ShapeDtypeStruct((n, n), jnp.float32))
    r_closure = float(exp_closure.call(jnp.asarray(X_np)))
    size_closure = len(exp_closure.serialize())
    print(f"闭包版：导出文件 {size_closure:,} 字节；结果 {r_closure:.4f}，参考 {ref:.4f}，误差 {abs(r_closure - ref):.2e}")
    print("闭包版是否触发了 CPU bug：", abs(r_closure - ref) > 1.0)
except Exception as e:
    print("实测报错：", type(e).__name__, e)
# 权重确实以 128×128 的常量形式嵌在 StableHLO 里
assert re.search(r"stablehlo\.constant dense<.*> : tensor<128x128xf32>", exp_closure.mlir_module())

# ✅ 正确写法：权重作为参数传入，导出的程序里不含任何权重
exp_param = export.export(jax.jit(lambda w, x: jnp.sum(w @ x)))(
    jax.ShapeDtypeStruct((n, n), jnp.float32), jax.ShapeDtypeStruct((n, n), jnp.float32))
r_param = float(exp_param.call(jnp.asarray(W_np), jnp.asarray(X_np)))
size_param = len(exp_param.serialize())
print(f"传参版：导出文件 {size_param:,} 字节；结果 {r_param:.4f}，误差 {abs(r_param - ref):.2e}")
print(f"闭包版文件是传参版的 {size_closure / size_param:.0f} 倍（权重 {W_np.nbytes:,} 字节被烘焙了进去）")

assert abs(r_param - ref) < 1e-2                                  # 传参版结果正确
assert size_closure > W_np.nbytes > 10 * size_param               # 闭包版的大小主要就是那份权重
assert not re.search(r"tensor<128x128xf32> = stablehlo\.constant|constant dense<[^>]{100,}", exp_param.mlir_module())

x_poly = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
print("符号形状:", x_poly.shape, "| 类型:", type(x_poly.shape[0]).__name__, "| is_symbolic_dim:", export.is_symbolic_dim(x_poly.shape[0]))

exp_poly = export.export(jax.jit(mlp_apply))(sds(params), x_poly)
print("in_avals :", exp_poly.in_avals)
print("out_avals:", exp_poly.out_avals)
print("@main 签名:", main_signature(exp_poly)[-120:])
assert str(exp_poly.out_avals[0].shape) == "(b, 3)"
assert "tensor<?x4xf32>" in main_signature(exp_poly) and "shape_assertion" in exp_poly.mlir_module()

exp_poly_back = export.deserialize(exp_poly.serialize())           # 符号形状也能序列化
for batch in [1, 5, 17, 256]:
    xb = jax.random.normal(jax.random.key(batch), (batch, 4))
    yb = exp_poly_back.call(params, xb)
    assert yb.shape == (batch, 3) and max_diff(yb, mlp_apply(params, xb)) < 1e-5
print("一份导出产物，batch = 1 / 5 / 17 / 256 全部与原函数一致 ✓")

# ❌ 错误写法：用不符合导出签名的形状 / dtype 调用
errs_call = {}
for name, bad_x in [("特征维 5 ≠ 4", jnp.ones((3, 5))), ("dtype int32", jnp.ones((3, 4), jnp.int32))]:
    try:
        exp_poly.call(params, bad_x)
        errs_call[name] = "可以运行"
    except ValueError as e:
        errs_call[name] = f"ValueError: {e}"
    print(f"{name}: {errs_call[name]}")
assert "Shape mismatch for args[1].shape[1]" in errs_call["特征维 5 ≠ 4"]
assert "Dtype mismatch for args[1]: expected float32 and called with int32" in errs_call["dtype int32"]

# ✅ 正确写法：在导出边界之外把输入整理成导出签名要求的形状和 dtype
raw = np.arange(12, dtype=np.int64).reshape(3, 4)            # 来自外部的数据可能是 int64
y_ok = exp_poly.call(params, jnp.asarray(raw, jnp.float32))
assert y_ok.shape == (3, 3)
print("输入先转成 float32、(b, 4) 再调用 ✓")

# ❌ 错误写法：reshape 成 (-1, 3)——4b 不一定能被 3 整除
try:
    export.export(jax.jit(lambda x: x.reshape(-1, 3)))(x_poly)
    print("实测：可以导出")
except Exception as e:
    err_div = f"{type(e).__name__}: {e}"
    print("实测报错：", err_div[:220])
assert err_div.startswith("InconclusiveDimensionOperation") and "Cannot divide evenly" in err_div

# ✅ 正确写法：只做对任意 b 都成立的 reshape；输出形状里出现符号表达式 2*b
exp_r2 = export.export(jax.jit(lambda x: x.reshape(-1, 2)))(x_poly)
print("reshape(-1, 2) 的输出形状:", exp_r2.out_avals[0].shape)
assert str(exp_r2.out_avals[0].shape[0]) == "2*b"
assert exp_r2.call(jnp.ones((5, 4))).shape == (10, 2)

# ❌ 错误写法：对符号维度做 Python 条件判断（b > 2 对有的 b 成立、有的不成立）
def scale_if_large(x):
    if x.shape[0] > 2:            # 追踪时要求一个确定的 Python bool
        return x * 2.0
    return x

try:
    export.export(jax.jit(scale_if_large))(x_poly)
    print("实测：可以导出")
except Exception as e:
    err_cmp = f"{type(e).__name__}: {e}"
    print("实测报错：", err_cmp.split("\n")[0])
assert err_cmp.startswith("InconclusiveDimensionOperation") and "'b' > '2' is inconclusive" in err_cmp

# ✅ 正确写法：用 constraints 声明 b >= 3，比较就有了确定结论；调用时约束会被检查
x_poly_ge3 = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4", constraints=("b >= 3",)), jnp.float32)
exp_ge3 = export.export(jax.jit(scale_if_large))(x_poly_ge3)
assert max_diff(exp_ge3.call(jnp.ones((5, 4))), 2.0 * jnp.ones((5, 4))) == 0.0
try:
    exp_ge3.call(jnp.ones((2, 4)))                       # b = 2 违反约束
except ValueError as e:
    err_ge3 = str(e)
    print("b = 2 调用时报错：", err_ge3[:110], "...")
assert "do not match the symbolic shape constraint b >= 3" in err_ge3
print("声明约束后可以导出；违反约束的调用被 shape_assertion 拦下 ✓")

# ❌ 错误写法：用 floordiv(b, 2) 隐含"b 是偶数"的假设——能导出，奇数 b 调用时在形状细化阶段失败
half_reshape = lambda x: x.reshape(x.shape[0] // 2, 2, 4)
exp_floor = export.export(jax.jit(half_reshape))(x_poly)
print("导出成功，输出形状:", exp_floor.out_avals[0].shape)
print("b = 4 调用:", exp_floor.call(jnp.ones((4, 4))).shape)
try:
    exp_floor.call(jnp.ones((3, 4)))
    print("b = 3：可以运行")
except ValueError as e:
    err_refine = str(e)
    print("b = 3 调用时报错：ValueError:", err_refine)
assert str(exp_floor.out_avals[0].shape[0]) == "floordiv(b, 2)"
assert "Error refining shapes" in err_refine

# ✅ 正确写法：把"偶数"写进形状规格 "2*k, 4"；b // 2 化简为 k，奇数输入在调用时被清楚地拒绝
x_even = jax.ShapeDtypeStruct(export.symbolic_shape("2*k, 4"), jnp.float32)
exp_even = export.export(jax.jit(half_reshape))(x_even)
print("输入规格:", exp_even.in_avals[0].shape, "→ 输出形状:", exp_even.out_avals[0].shape)
assert exp_even.call(jnp.ones((6, 4))).shape == (3, 2, 4)
try:
    exp_even.call(jnp.ones((5, 4)))
except ValueError as e:
    err_even = str(e)
    print("b = 5 调用时报错：", err_even[:150], "...")
assert str(exp_even.out_avals[0].shape[0]) == "k"
assert "Division had remainder 1 when computing the value of 'k'" in err_even

mask_fn = lambda x, m: x * m[:, None]

# ❌ 错误写法：两次独立调用 symbolic_shape，两个 "b" 属于不同的作用域
try:
    spec_x = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
    spec_m = jax.ShapeDtypeStruct(export.symbolic_shape("b"), jnp.float32)
    export.export(jax.jit(mask_fn))(spec_x, spec_m)
    print("实测：可以导出")
except ValueError as e:
    err_scope = str(e)
    print("实测报错：ValueError:", err_scope.split("\n")[0])
assert "Invalid mixing of symbolic scopes" in err_scope

# ✅ 正确写法：symbolic_args_specs 用同一个作用域生成所有参数的规格（"_" 表示沿用示例参数的维度）
specs = export.symbolic_args_specs((jnp.ones((2, 4)), jnp.ones((2,))), ("b, _", "b"))
print("生成的规格:", specs)
exp_mask = export.export(jax.jit(mask_fn))(*specs)
xm, mm = jnp.arange(12.).reshape(3, 4), jnp.array([1., 0., 2.])
assert max_diff(exp_mask.call(xm, mm), mask_fn(xm, mm)) == 0.0
try:
    exp_mask.call(jnp.ones((3, 4)), jnp.ones(5))           # 两个参数的 b 不一致
except ValueError as e:
    print("b 不一致时报错：", str(e)[:120], "...")
    assert "Found inconsistency between dimension size args[1].shape[0] (= 5)" in str(e)

# 实测上面的第一、二条
try:
    a_, b_ = export.symbolic_shape("a, b")
    export.export(jax.jit(lambda x: x))(jax.ShapeDtypeStruct((a_ + b_, 4), jnp.float32))
except ValueError as e:
    print("规格 a+b：", str(e)[:105], "...")
    assert "Cannot solve for values of dimension variables" in str(e)

b_sym, = export.symbolic_shape("b")
print("b == 1 →", b_sym == 1, "| b >= 1 →", b_sym >= 1)
assert (b_sym == 1) is False and (b_sym >= 1) is True

# 只做 lowering，未在 GPU/TPU 执行
exp_cpu = export.export(jax.jit(mlp_apply))(sds(params), x_poly)
exp_mp = export.export(jax.jit(mlp_apply), platforms=("cpu", "cuda", "tpu"))(sds(params), x_poly)
print("platforms:", exp_mp.platforms)
print("多平台 @main 签名开头:", main_signature(exp_mp)[:110], "...")
size_cpu, size_mp = len(exp_cpu.serialize()), len(exp_mp.serialize())
print(f"序列化大小：仅 cpu {size_cpu:,} 字节；cpu+cuda+tpu {size_mp:,} 字节（多 {size_mp / size_cpu - 1:.0%}）")

assert exp_mp.platforms == ("cpu", "cuda", "tpu")
assert "platform_index" in main_signature(exp_mp) and "platform_index" not in main_signature(exp_cpu)
assert size_mp < 2 * size_cpu                        # 共享的部分只存一份
# 在本机（CPU）调用：自动选中 cpu 分支，结果与原函数一致
xb = jax.random.normal(jax.random.key(7), (6, 4))
assert max_diff(export.deserialize(exp_mp.serialize()).call(params, xb), mlp_apply(params, xb)) < 1e-5
print("多平台产物在 CPU 上调用（自动选择 cpu 分支）结果正确 ✓")

# ❌ 错误写法：只为 cuda 导出，却拿到 CPU 上调用
exp_cuda = export.export(jax.jit(mlp_apply), platforms=("cuda",))(sds(params), x_poly)   # 只做 lowering
try:
    exp_cuda.call(params, xb)
    print("实测：可以运行")
except ValueError as e:
    err_plat = str(e)
    print("实测报错：ValueError:", err_plat)
assert "was exported for platforms '('cuda',)' but it is used on '('cpu',)'" in err_plat

# ✅ 正确写法：platforms 里包含所有会运行它的平台（包括本地开发 / 测试用的 cpu）
exp_cuda_cpu = export.export(jax.jit(mlp_apply), platforms=("cuda", "cpu"))(sds(params), x_poly)
assert max_diff(exp_cuda_cpu.call(params, xb), mlp_apply(params, xb)) < 1e-5
print("platforms=('cuda', 'cpu')：CPU 上可直接调用；cuda 分支只做了 lowering ✓")

import inspect
print("Exported.serialize 签名:", inspect.signature(export.Exported.serialize))
assert "vjp_order" in str(inspect.signature(export.Exported.serialize))

def mse_loss(params, x, y):
    return jnp.mean((mlp_apply(params, x) - y) ** 2)

xg = jax.random.normal(jax.random.key(10), (16, 4))
yg = jax.random.normal(jax.random.key(11), (16, 3))
exp_loss = export.export(jax.jit(mse_loss))(sds(params), sds(xg), sds(yg))
print("同一进程内 has_vjp():", exp_loss.has_vjp())

g_ref = jax.grad(mse_loss)(params, xg, yg)
g_exp = jax.grad(exp_loss.call)(params, xg, yg)            # 对导出函数求导，梯度仍是字典
print("梯度结构:", jax.tree.map(jnp.shape, g_exp))
assert exp_loss.has_vjp() and max_diff(g_exp, g_ref) < 1e-6
vjp_exp = exp_loss.vjp()                                    # 反向函数本身也是一个 Exported
print("vjp() 的输入:", len(vjp_exp.in_avals), "个（原输入 6 个叶子 + 输出余切 1 个）；输出:", len(vjp_exp.out_avals), "个")
assert len(vjp_exp.in_avals) == 7 and len(vjp_exp.out_avals) == 6

# ❌ 错误写法：用默认的 serialize()（vjp_order=0）保存，读回后再求导
loaded0 = export.deserialize(exp_loss.serialize())
try:
    jax.grad(loaded0.call)(params, xg, yg)
    print("实测：可以求导")
except ValueError as e:
    err_vjp0 = str(e)
    print("实测报错：ValueError:", err_vjp0, "| has_vjp() =", loaded0.has_vjp())
assert err_vjp0 == "No VJP is available" and not loaded0.has_vjp()

# ✅ 正确写法：导出端声明需要的阶数 vjp_order=k，k 阶以内的反向导数都可用
sizes = {k: len(exp_loss.serialize(vjp_order=k)) for k in (0, 1, 2)}
print("序列化大小（字节）:", sizes)
assert sizes[0] < sizes[1] < sizes[2]

loaded1 = export.deserialize(exp_loss.serialize(vjp_order=1))
assert max_diff(jax.grad(loaded1.call)(params, xg, yg), g_ref) < 1e-6           # 一阶 OK

def grad_norm2(load):                                       # 二阶：对"梯度范数平方"再求导
    return lambda p: sum(jnp.sum(g ** 2) for g in jax.tree.leaves(jax.grad(load)(p, xg, yg)))
try:
    jax.grad(grad_norm2(loaded1.call))(params)
except ValueError as e:
    print("vjp_order=1 时求二阶导：ValueError:", e)
    assert str(e) == "No VJP is available"

loaded2 = export.deserialize(exp_loss.serialize(vjp_order=2))
gg_ref = jax.grad(grad_norm2(mse_loss))(params)
gg_exp = jax.grad(grad_norm2(loaded2.call))(params)
assert max_diff(gg_exp, gg_ref) < 1e-5
print("vjp_order=1：一阶导正确、二阶导报错；vjp_order=2：二阶导与原函数一致 ✓")

# ❌ 错误写法：对 Exported.call 做 vmap / jvp
x_single = jax.random.normal(jax.random.key(12), (2, 4))
x_stack = jax.random.normal(jax.random.key(13), (5, 2, 4))      # 5 组，每组 (2, 4)
errs_tf = {}
for name, fn in [("vmap", lambda: jax.vmap(exp.call, in_axes=(None, 0))(params, x_stack)),
                 ("jvp", lambda: jax.jvp(lambda x: exp.call(params, x), (x_single,), (x_single,)))]:
    try:
        fn()
        errs_tf[name] = "可以运行"
    except Exception as e:
        errs_tf[name] = f"{type(e).__name__}: {e}"
    print(f"{name}: {errs_tf[name]}")
assert errs_tf["vmap"].startswith("NotImplementedError: Batching rule for 'call_exported' not implemented")
assert errs_tf["jvp"].startswith("TypeError") and "custom_vjp" in errs_tf["jvp"]

# ✅ 正确写法：在导出前把批维度做进程序——导出 vmap 之后的函数（这里再加一个符号维度 n）
batched = jax.vmap(mlp_apply, in_axes=(None, 0))
x_stack_spec = jax.ShapeDtypeStruct(export.symbolic_shape("n, 2, 4"), jnp.float32)
exp_batched = export.export(jax.jit(batched))(sds(params), x_stack_spec)
assert max_diff(exp_batched.call(params, x_stack), batched(params, x_stack)) < 1e-5

# jit 可以嵌套：导出的函数可以作为更大 jit 函数的一部分（前后处理用 jnp 写）
@jax.jit
def predict_labels(params, x):
    return jnp.argmax(exp_poly.call(params, x), axis=-1)
assert bool(jnp.all(predict_labels(params, xb) == jnp.argmax(mlp_apply(params, xb), -1)))
print("导出 vmap(f) 代替 vmap(exp.call) ✓；Exported.call 可以嵌进外层 jit ✓")
jaxpr_txt = str(jax.make_jaxpr(predict_labels)(params, xb))
print("外层 jaxpr 里含 call_exported 原语：", "call_exported" in jaxpr_txt)
assert "call_exported" in jaxpr_txt

import jax._src.export.serialization as _ser            # 私有模块，这里只读取一个常量来查看，不要在代码里依赖它
from jax.extend.mlir.dialects import stablehlo as shlo

print("序列化格式版本 _SERIALIZATION_VERSION =", _ser._SERIALIZATION_VERSION)
print("调用约定：默认", jax.config.jax_export_calling_convention_version,
      "| 支持范围", export.minimum_supported_calling_convention_version, "–", export.maximum_supported_calling_convention_version)
target = shlo.get_version_from_compatibility_requirement(shlo.StablehloCompatibilityRequirement.WEEK_4)
print("StableHLO：当前版本", shlo.get_current_version(), "| 最低可读版本", shlo.get_minimum_version(), "| 序列化目标版本（≥4 周前）", target)
bc = exp.mlir_module_serialized
print("可移植字节码文件头:", bc[:4], "（MLIR bytecode 魔数 'ML\\xefR'）")

assert isinstance(_ser._SERIALIZATION_VERSION, int) and _ser._SERIALIZATION_VERSION >= 12
assert export.minimum_supported_calling_convention_version <= exp.calling_convention_version <= export.maximum_supported_calling_convention_version
assert exp.calling_convention_version == jax.config.jax_export_calling_convention_version == 10
ver = lambda s: tuple(int(t) for t in s.split("."))
assert ver(shlo.get_minimum_version()) <= ver(target) <= ver(shlo.get_current_version())
assert bc[:4] == b"ML\xefR"

cc9_path = SCRATCH / "mlp_cc9.jaxexport"
out = run_py("""
    import sys, jax, jax.numpy as jnp
    from jax import export
    w = jax.ShapeDtypeStruct((4, 3), jnp.float32)
    x = jax.ShapeDtypeStruct(export.symbolic_shape("b, 4"), jnp.float32)
    e = export.export(jax.jit(lambda w, x: jnp.tanh(x @ w)))(w, x)
    open(sys.argv[1], "wb").write(e.serialize())
    print("CC", e.calling_convention_version)
""", cc9_path, env_extra={"JAX_EXPORT_CALLING_CONVENTION_VERSION": "9"})
print("子进程输出:", out.strip())
e9 = export.deserialize(cc9_path.read_bytes())
w9 = jax.random.normal(jax.random.key(3), (4, 3))
x9 = jax.random.normal(jax.random.key(4), (5, 4))
print("当前进程读回：calling_convention_version =", e9.calling_convention_version)
assert out.strip() == "CC 9" and e9.calling_convention_version == 9
assert max_diff(e9.call(w9, x9), jnp.tanh(x9 @ w9)) < 1e-6
print("版本 9 的产物在默认版本 10 的环境里正常运行 ✓")

# 实测：Pallas TPU kernel 不需要 DisabledSafetyCheck，自定义 FFI 目标需要（只做 lowering，未在 TPU 上执行）
from jax.experimental import pallas as pl
from jax._src.export._export import _CUSTOM_CALL_TARGETS_GUARANTEED_STABLE as _STABLE   # 私有常量，这里只读取查看

def _add_kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]
s8 = jax.ShapeDtypeStruct((8, 128), jnp.float32)
exp_pl = export.export(jax.jit(pl.pallas_call(_add_kernel, out_shape=s8)), platforms=("tpu",))(s8, s8)
blob_pl = exp_pl.serialize()                                     # 没有传 disabled_checks
assert "tpu_custom_call" in exp_pl.mlir_module() and len(blob_pl) > 0
print("Pallas TPU kernel：导出 + 序列化成功（只做 lowering）")

ffi_fn = lambda x: jax.ffi.ffi_call("my_custom_target", jax.ShapeDtypeStruct(x.shape, x.dtype))(x)
try:
    export.export(jax.jit(ffi_fn))(jnp.ones(3))
    err_ffi = None
except ValueError as e:
    err_ffi = str(e)
    print("自定义 FFI 目标：ValueError:", err_ffi.splitlines()[0][:95], "...")
assert err_ffi is not None and "no compatibility guarantees" in err_ffi
exp_ffi = export.export(jax.jit(ffi_fn),
                        disabled_checks=[export.DisabledSafetyCheck.custom_call("my_custom_target")])(jnp.ones(3))
assert len(exp_ffi.serialize()) > 0
print("显式 DisabledSafetyCheck.custom_call 后可以导出（消费端必须注册同名目标才能运行）")
print("稳定名单：tpu_custom_call", "tpu_custom_call" in _STABLE, "| mosaic_gpu_v2", "mosaic_gpu_v2" in _STABLE,
      "| __gpu$xla.gpu.triton", "__gpu$xla.gpu.triton" in _STABLE)
assert {"tpu_custom_call", "mosaic_gpu_v2"} <= _STABLE and "__gpu$xla.gpu.triton" not in _STABLE

shard_src = """
    import os, sys
    import numpy as np, jax, jax.numpy as jnp
    from jax import export, lax
    from jax.sharding import PartitionSpec as P
    assert jax.device_count() == 4
    def check(name, ok, info=""):
        print(("PASS " if ok else "FAIL ") + name, info)

    f = lambda w, x: jnp.tanh(x @ w)
    w = jnp.full((4, 3), 0.1); x = jnp.arange(32.).reshape(8, 4)
    mesh = jax.make_mesh((4,), ("data",))
    with jax.set_mesh(mesh):
        xs, ws = jax.device_put(x, P("data")), jax.device_put(w, P())
        e = export.export(jax.jit(f))(ws, xs)                    # 输入按 batch 维分到 4 个设备
        check("1 nr_devices==4", e.nr_devices == 4, e.nr_devices)
        check("1 记录输入分片", str(e.in_shardings_jax(mesh)[1].spec) == str(P("data")), e.in_shardings_jax(mesh)[1].spec)
        y = e.call(ws, xs)
        check("1 同 mesh 调用正确", np.allclose(y, f(w, x), atol=1e-6), y.sharding.spec)
    blob = e.serialize()
    e2 = export.deserialize(blob)
    try:                                                         # 不在 mesh 里调用
        e2.call(w, x); check("1 无 mesh 调用应失败", False)
    except ValueError as err:
        check("1 无 mesh 调用 → ValueError", "is not equal to the size of the mesh 4" in str(err), str(err)[:90])
    mesh2 = jax.make_mesh((2,), ("data",), devices=jax.devices()[:2])
    try:                                                         # 只用 2 个设备的 mesh
        with jax.set_mesh(mesh2):
            e2.call(jax.device_put(w, P()), jax.device_put(x, P("data")))
        check("1 2 设备调用应失败", False)
    except ValueError as err:
        check("1 2 设备 mesh 调用 → ValueError", "Length of device assignment 2" in str(err), str(err)[:90])
    mesh_rev = jax.make_mesh((4,), ("data",), devices=jax.devices()[::-1])
    with jax.set_mesh(mesh_rev):                                 # 同样 4 个设备、顺序不同
        y = e2.call(jax.device_put(w, P()), jax.device_put(x, P("data")))
        check("1 换一组 4 设备 mesh 调用正确", np.allclose(y, f(w, x), atol=1e-6))

    e1 = export.export(jax.jit(f))(w, x)                          # 规则 2：单设备导出
    with jax.set_mesh(mesh):
        y = e1.call(ws, xs)
        check("2 单设备产物接受分片输入", e1.nr_devices == 1 and np.allclose(y, f(w, x), atol=1e-6), y.sharding.spec)

    g = lambda v: lax.psum(jnp.sum(v), "data")                   # 规则 3：shard_map
    with jax.set_mesh(mesh):
        sm = jax.jit(jax.shard_map(g, in_specs=P("data"), out_specs=P()))
        e3 = export.export(sm)(xs)
        txt = e3.mlir_module()
        check("3 shard_map 可导出且含 all_reduce", e3.nr_devices == 4 and "all_reduce" in txt)
        check("3 shard_map 结果正确", float(e3.call(xs)) == float(x.sum()), float(e3.call(xs)))

    try:                                                          # 规则 4：pmap
        export.export(jax.pmap(lambda v: v * 2))(jnp.ones((4, 3))); check("4 直接导出 pmap 应失败", False)
    except ValueError as err:
        check("4 直接导出 pmap → ValueError", "must be the result of `jit`" in str(err))
    ep = export.export(jax.jit(jax.pmap(lambda v: v * 2)))(jnp.ones((4, 3)))
    check("4 jit(pmap) 可导出", ep.nr_devices == 4, ep.nr_devices)
"""
out = run_py(shard_src, env_extra={"XLA_FLAGS": "--xla_force_host_platform_device_count=4"})
print(out)
lines = [l for l in out.splitlines() if l.startswith(("PASS", "FAIL"))]
assert len(lines) == 11 and all(l.startswith("PASS") for l in lines), lines

ckpt_path = SCRATCH / "single_ckpt"
state0 = {"w": jnp.arange(6.).reshape(2, 3), "step": jnp.array(0)}
ckptr.save(ckpt_path, state0)
ckptr.wait_until_finished()

# ❌ 错误写法：对同一个目录再次 save（比如每个 epoch 覆盖同一个"最新检查点"目录）
try:
    ckptr.save(ckpt_path, {"w": state0["w"] + 1, "step": jnp.array(1)})
    ckptr.wait_until_finished()
    print("实测：可以覆盖")
except ValueError as e:
    err_exists = str(e)
    print("实测报错：ValueError:", err_exists.replace(str(SCRATCH), "<SCRATCH>"))
assert "already exists" in err_exists

# ✅ 正确写法：确实要覆盖时显式传 force=True；需要多步保存时用 CheckpointManager（§3.2），每步一个子目录
state1 = {"w": state0["w"] + 1, "step": jnp.array(1)}
ckptr.save(ckpt_path, state1, force=True)
ckptr.wait_until_finished()
restored1 = ckptr.restore(ckpt_path, sds(state1))
assert int(restored1["step"]) == 1 and max_diff(restored1, state1) == 0.0
print("force=True 覆盖成功，读回的是新内容 ✓")

mgr_dir = SCRATCH / "mgr_interval"
mngr = ocp.CheckpointManager(mgr_dir, options=ocp.CheckpointManagerOptions(max_to_keep=2, save_interval_steps=2))
saved_flags = [mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step})) for step in range(7)]
mngr.wait_until_finished()
print("每步 save() 的返回值:", saved_flags)
print("磁盘上保留的步:", mngr.all_steps(), "| 子目录:", sorted(p.name for p in mgr_dir.iterdir()))
print("latest_step:", mngr.latest_step())
assert saved_flags == [True, False, True, False, True, False, True]
assert mngr.all_steps() == [4, 6] and mngr.latest_step() == 6
restored = mngr.restore(6, args=ocp.args.StandardRestore(jax.eval_shape(lambda: {"w": jnp.ones(3)})))
assert max_diff(restored, {"w": jnp.ones(3) * 6}) == 0.0
mngr.close()

losses = [5.0, 1.0, 4.0, 2.0, 9.0]          # 第 4 步（最新）的验证损失最差

# ❌ 错误写法：只靠 best_fn + max_to_keep，以为会"保留最好的 2 个 + 最新的"
try:
    mngr = ocp.CheckpointManager(SCRATCH / "mgr_best", options=ocp.CheckpointManagerOptions(
        max_to_keep=2, best_fn=lambda m: m["loss"], best_mode="min"))
    for step, l in enumerate(losses):
        mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step}), metrics={"loss": l})
    mngr.wait_until_finished()
    print("实测：保留的步", mngr.all_steps(), "| best_step", mngr.best_step(), "| latest_step", mngr.latest_step())
    kept_bad, latest_bad = mngr.all_steps(), mngr.latest_step()
    mngr.close()
except Exception as e:
    print("实测报错：", type(e).__name__, e)
assert kept_bad == [1, 3] and 4 not in kept_bad          # 最新的第 4 步被删掉了
assert latest_bad == 3                                   # latest_step() 也"倒退"到了第 3 步

# ✅ 正确写法：用 preservation_policy 组合"最新 1 个"和"最好 2 个"
cm = ocp.checkpoint_managers
policy = cm.AnyPreservationPolicy([
    cm.LatestN(n=1),                                                    # 断点续训用
    cm.BestN(get_metric_fn=lambda m: m["loss"], reverse=True, n=2),     # reverse=True：指标越小越好（实测）
])
mngr = ocp.CheckpointManager(SCRATCH / "mgr_policy", options=ocp.CheckpointManagerOptions(
    preservation_policy=policy, best_fn=lambda m: m["loss"], best_mode="min"))
for step, l in enumerate(losses):
    mngr.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3) * step}), metrics={"loss": l})
mngr.wait_until_finished()
print("保留的步:", mngr.all_steps(), "| best_step:", mngr.best_step(), "| latest_step:", mngr.latest_step())
assert mngr.all_steps() == [1, 3, 4] and mngr.best_step() == 1 and mngr.latest_step() == 4

best = mngr.restore(mngr.best_step(), args=ocp.args.StandardRestore(jax.eval_shape(lambda: {"w": jnp.ones(3)})))
assert max_diff(best, {"w": jnp.ones(3) * 1}) == 0.0
mngr.close()
print("最新 1 个 + 最好 2 个都在；按 best_step 恢复得到第 1 步的权重 ✓")

kept = {}
for rev in (False, True):
    m_ = ocp.CheckpointManager(SCRATCH / f"mgr_rev_{rev}", options=ocp.CheckpointManagerOptions(
        preservation_policy=cm.BestN(get_metric_fn=lambda m: m["loss"], reverse=rev, n=2)))
    for step, l in enumerate(losses):
        m_.save(step, args=ocp.args.StandardSave({"w": jnp.ones(3)}), metrics={"loss": l})
    m_.wait_until_finished()
    kept[rev] = m_.all_steps()
    m_.close()
print("losses =", losses, "→ reverse=False 保留", kept[False], "；reverse=True 保留", kept[True])
assert kept[False] == [0, 4] and kept[True] == [1, 3]

init_fn = lambda key: init_mlp(key)
abstract = jax.eval_shape(init_fn, jax.random.key(0))        # 不分配内存
print("抽象模板:", jax.tree.map(lambda s: f"{s.dtype}{list(s.shape)}", abstract))
ckptr.save(SCRATCH / "mlp_params", params)
ckptr.wait_until_finished()

restored_f32 = ckptr.restore(SCRATCH / "mlp_params", abstract)
assert max_diff(restored_f32, params) == 0.0

abstract_bf16 = jax.tree.map(lambda s: jax.ShapeDtypeStruct(s.shape, jnp.bfloat16), abstract)
restored_bf16 = ckptr.restore(SCRATCH / "mlp_params", abstract_bf16)
print("按 bf16 模板恢复:", {k: str(v.dtype) for k, v in restored_bf16.items()})
assert all(v.dtype == jnp.bfloat16 for v in restored_bf16.values())
assert max_diff(restored_bf16, params) < 0.02
print("用 eval_shape 模板恢复 fp32 逐位一致；按 bf16 模板恢复时自动转换 ✓")

# ❌ 错误写法：模板的形状与磁盘上的不一致（比如改了隐藏层宽度却加载旧检查点）
wrong = dict(abstract, w1=jax.ShapeDtypeStruct((4, 16), jnp.float32))
try:
    ckptr.restore(SCRATCH / "mlp_params", wrong)
    print("实测：可以恢复")
except ValueError as e:
    err_shape = str(e)
    print("实测报错：ValueError:", err_shape[:150], "...")
assert "is not compatible with the stored shape: (4, 8)" in err_shape

# ✅ 正确写法：先用 metadata() 查看磁盘上的结构和形状，再决定模板（或者确认要加载的是哪个版本的模型）
meta_tree = ckptr.metadata(SCRATCH / "mlp_params").item_metadata.tree
print("磁盘上的形状:", {k: tuple(v.shape) for k, v in meta_tree.items()})
assert {k: tuple(v.shape) for k, v in meta_tree.items()} == {k: v.shape for k, v in params.items()}

class MLP(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs: nnx.Rngs):
        self.fc1 = nnx.Linear(din, dh, rngs=rngs)
        self.fc2 = nnx.Linear(dh, dh, rngs=rngs)
        self.fc3 = nnx.Linear(dh, dout, rngs=rngs)

    def __call__(self, x):
        x = nnx.relu(self.fc1(x))
        x = nnx.relu(self.fc2(x))
        return self.fc3(x)

TEACHER = jax.random.normal(jax.random.key(100), (4, 3))      # 生成标签用的"教师"矩阵

def make_batch(data_key, step, teacher, n=128):
    """第 step 步的数据批次：由 (data_key, step) 决定，恢复了 step 就恢复了数据位置。"""
    x = jax.random.normal(jax.random.fold_in(data_key, step), (n, 4))
    y = jnp.argmax(jnp.sin(2.0 * x) @ teacher, axis=-1)
    return x, y

def loss_fn(model, x, y):
    return optax.softmax_cross_entropy_with_integer_labels(model(x), y).mean()

@nnx.jit
def train_step(model, optimizer, data_key, teacher):
    x, y = make_batch(data_key, optimizer.step[...], teacher)          # 数据位置取自优化器步数
    loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)
    optimizer.update(model, grads)
    return loss

TX = optax.adam(1e-2)
def make_model_and_opt():
    model = MLP(4, 32, 3, rngs=nnx.Rngs(0))
    return model, nnx.Optimizer(model, TX, wrt=nnx.Param)

DATA_KEY = jax.random.key(42)
N1, N2 = 30, 30

# (a) 一口气训练 N1 + N2 步，作为参照
model_ref, opt_ref = make_model_and_opt()
curve_ref = [float(train_step(model_ref, opt_ref, DATA_KEY, TEACHER)) for _ in range(N1 + N2)]

# (b) 训练 N1 步后保存
model_a, opt_a = make_model_and_opt()
curve_a = [float(train_step(model_a, opt_a, DATA_KEY, TEACHER)) for _ in range(N1)]
nnx_dir = SCRATCH / "nnx_ckpt"
mngr = ocp.CheckpointManager(nnx_dir, options=ocp.CheckpointManagerOptions(max_to_keep=3))
mngr.save(int(opt_a.step[...]), args=ocp.args.StandardSave({"model": nnx.state(model_a), "opt": nnx.state(opt_a)}),
          metrics={"loss": curve_a[-1]})
mngr.wait_until_finished()
print("已保存的步:", mngr.all_steps(), "| State 类型:", type(nnx.state(model_a)).__name__)
assert mngr.all_steps() == [N1]
del model_a, opt_a                                   # 模拟进程退出：内存里的模型没了