Blame for tests/make_golden.py

"""生成 ZeroPoint 的参考数据 tests/golden/zeropoint_ref.json：固定的参数 + 输入 + 标准输出。

所有实现（zerocg 并行形式、numpy 递推形式、以后内核里的 C 版本）都拿它对答案，
谁和它对不上谁就有 bug。**只有在有意修改架构时才重新生成**，并在提交说明里写清楚为什么。

    python3 tests/make_golden.py
"""
import json
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
import zerocg as zc  # noqa: E402
from models import ZeroX  # noqa: E402

CONFIG = dict(vocab_size=13, n_embd=16, n_head=2, n_slot=4, n_layer=2)

if __name__ == "__main__":
    zc.set_default_dtype(np.float64)
    np.random.seed(2026)
    model = ZeroX(**CONFIG).eval()
    # 参数不用初始值：随机扰动一下，避免初始化时的对称性掩盖 bug
    for p in model.parameters():
        p.data = p.data + np.random.randn(*p.shape) * 0.3
    idx = np.random.randint(0, CONFIG["vocab_size"], size=(2, 70))   # 跨 3 个块
    logits = model(idx).data
    out = os.path.join(HERE, "golden", "zeropoint_ref.json")
    data = {"idx": idx.tolist(), "logits": logits.tolist()}
    data.update({"param:" + k: v.tolist() for k, v in model.state_dict().items()})
    data.update({"config:" + k: int(v) for k, v in CONFIG.items()})
    with open(out, "w", encoding="utf-8") as f:
        json.dump(data, f, separators=(",", ":"), ensure_ascii=False)
        f.write("\n")
    print(f"写入 {out}：输入 {idx.shape}，输出 {logits.shape}，{len(model.state_dict())} 个参数张量")