Blame for tests/test_zerocg.py

"""框架自检：用数值差分验证每个算子 / 层的反向传播，并跑两个小模型确认能学会。
运行: python3 tests/test_zerocg.py
"""
import os
import sys

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import zerocg as zc  # noqa: E402
from zerocg import functional as F  # noqa: E402
from zerocg import nn  # noqa: E402

zc.set_default_dtype(np.float64)   # 数值梯度检查需要 float64 的精度
rng = np.random.default_rng(0)


def grad_check(fn, *shapes, tol=1e-5, positive=False):
    """fn(*tensors) -> Tensor。比较自动求导和中心差分的梯度（float64）。"""
    arrays = [rng.standard_normal(s) for s in shapes]
    if positive:
        arrays = [np.abs(a) + 0.5 for a in arrays]
    ts = [zc.Tensor(a.copy(), requires_grad=True) for a in arrays]
    w = np.asarray(rng.standard_normal(fn(*ts).shape))  # 随机投影成标量，检查完整雅可比
    (fn(*ts) * zc.Tensor(w)).sum().backward()
    eps = 1e-6
    for t, a in zip(ts, arrays):
        num = np.zeros_like(a)
        for i in np.ndindex(a.shape):
            for sign in (1, -1):
                b = a.copy()
                b[i] += sign * eps
                inputs = [zc.Tensor(b) if x is t else zc.Tensor(x.data) for x in ts]
                with zc.no_grad():
                    num[i] += sign * (fn(*inputs).data * w).sum()
            num[i] /= 2 * eps
        err = np.abs(num - t.grad).max() / max(1.0, np.abs(num).max())
        assert err < tol, f"梯度误差 {err:.2e}"


def module_grad_check(module, x):
    for p in module.parameters():
        p.data = p.data.astype(np.float64)
    params = module.parameters()
    # 直接对参数逐个做差分
    xt = zc.Tensor(x)
    w = np.asarray(rng.standard_normal(module(xt).shape))
    for p in params:
        p.grad = None
    (module(xt) * zc.Tensor(w)).sum().backward()
    eps = 1e-6
    for p in params:
        num = np.zeros_like(p.data)
        for i in np.ndindex(p.shape):
            old = p.data[i]
            p.data[i] = old + eps
            with zc.no_grad():
                up = (module(xt).data * w).sum()
            p.data[i] = old - eps
            with zc.no_grad():
                down = (module(xt).data * w).sum()
            p.data[i] = old
            num[i] = (up - down) / (2 * eps)
        err = np.abs(num - p.grad).max() / max(1.0, np.abs(num).max())
        assert err < 1e-5, f"参数梯度误差 {err:.2e}"


# ---------------- 算子 ----------------
def test_ops():
    grad_check(lambda a, b: a + b, (3, 4), (4,))
    grad_check(lambda a, b: a - b, (3, 4), (3, 1))
    grad_check(lambda a, b: a * b, (2, 3), (2, 3))
    grad_check(lambda a, b: a / b, (2, 3), (3,), positive=True)
    grad_check(lambda a: a ** 3, (4,))
    grad_check(lambda a: a.sqrt(), (4,), positive=True)
    grad_check(lambda a: -a + 2.0 - a * 3.0, (4,))
    grad_check(lambda a: 1.0 / a, (4,), positive=True)
    grad_check(lambda a: a.exp(), (3, 2))
    grad_check(lambda a: a.log(), (3, 2), positive=True)
    grad_check(lambda a: a.tanh(), (5,))
    grad_check(lambda a: a.relu(), (5,))
    grad_check(lambda a: a.masked_fill(np.eye(3, dtype=bool), -5.0), (3, 3))


def test_reductions_and_shapes():
    grad_check(lambda a: a.sum(), (3, 4))
    grad_check(lambda a: a.sum(axis=1), (3, 4))
    grad_check(lambda a: a.mean(axis=0, keepdims=True), (3, 4))
    grad_check(lambda a: a.max(axis=-1), (3, 4))
    grad_check(lambda a: a.reshape(4, 3), (3, 4))
    grad_check(lambda a: a.transpose(2, 0, 1), (2, 3, 4))
    grad_check(lambda a: a[1:, ::2], (3, 4))
    grad_check(lambda a: a[np.array([0, 2, 0])], (3, 4))  # 重复索引
    grad_check(lambda a, b: zc.concat([a, b, a], axis=1), (2, 3), (2, 1))


def test_matmul():
    grad_check(lambda a, b: a @ b, (3, 4), (4, 5))
    grad_check(lambda a, b: a @ b, (2, 3, 4), (4, 5))      # 批量 @ 共享权重
    grad_check(lambda a, b: a @ b, (2, 2, 3, 4), (2, 2, 4, 3))


def test_functional():
    grad_check(lambda a: F.softmax(a, axis=-1), (3, 5))
    grad_check(lambda a: F.log_softmax(a, axis=-1), (3, 5))
    grad_check(lambda a: F.gelu(a), (6,))
    grad_check(lambda a: F.sigmoid(a), (6,))
    grad_check(lambda a: F.cross_entropy(a, np.array([1, 0, 4])), (3, 5))
    # 和直接的数学定义对比数值
    x = rng.standard_normal((4, 6))
    t = np.array([0, 5, 2, 2])
    ref = -np.log(np.exp(x) / np.exp(x).sum(1, keepdims=True))[np.arange(4), t].mean()
    assert abs(F.cross_entropy(zc.Tensor(x), t).item() - ref) < 1e-10


def test_layers():
    module_grad_check(nn.Linear(4, 3), rng.standard_normal((2, 4)))
    module_grad_check(nn.LayerNorm(5), rng.standard_normal((3, 5)))
    module_grad_check(nn.Sequential(nn.Linear(3, 4), nn.GELU(), nn.Linear(4, 2)), rng.standard_normal((2, 3)))
    module_grad_check(nn.CausalSelfAttention(8, 2), rng.standard_normal((2, 4, 8)))

    emb = nn.Embedding(5, 3)
    out = emb(np.array([[1, 1, 4]]))
    out.sum().backward()
    assert emb.weight.grad[1].tolist() == [2, 2, 2] and emb.weight.grad[0].tolist() == [0, 0, 0]


def test_attention_is_causal():
    attn = nn.CausalSelfAttention(8, 2).eval()
    x = rng.standard_normal((1, 5, 8)).astype(np.float32)
    y1 = attn(zc.Tensor(x)).data
    x[0, 4] += 10  # 改最后一个位置，前面的输出不应变化
    y2 = attn(zc.Tensor(x)).data
    assert np.allclose(y1[0, :4], y2[0, :4]) and not np.allclose(y1[0, 4], y2[0, 4])


def test_count_flops():
    with zc.count_flops() as c:
        zc.randn(2, 3, 4) @ zc.randn(4, 5)
    assert c.flops == 2 * (2 * 3 * 5) * 4


def test_default_dtype():
    zc.set_default_dtype(np.float32)
    try:
        x = zc.Tensor(np.ones(3))                       # float64 数组进来也变成 float32
        assert x.dtype == np.float32 and (x * 0.5).dtype == np.float32
        assert nn.Linear(2, 2).weight.dtype == np.float32
        assert zc.Tensor(np.arange(3)).dtype.kind == "i"   # 整数下标不动
        assert np.isfinite(F.sigmoid(zc.Tensor([-1000.0, 0.0, 1000.0])).data).all()
    finally:
        zc.set_default_dtype(np.float64)


def test_state_dict_roundtrip(tmp="/tmp"):
    path = os.path.join(os.environ.get("TMPDIR", tmp), "zerocg_test.npz")
    a = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4))
    a.save(path)
    b = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)).load(path)
    x = zc.randn(2, 3)
    assert np.allclose(a(x).data, b(x).data)
    assert [n for n, _ in a.named_parameters()] == ["layers.0.weight", "layers.0.bias", "layers.1.weight", "layers.1.bias"]


def test_optimizer_resume():
    """训练 10 步 == 训练 5 步、存档、新建模型和优化器、读档、再训练 5 步。"""
    X, y = rng.standard_normal((8, 3)), np.array([0, 1, 1, 0, 1, 0, 0, 1])

    def make():
        np.random.seed(7)
        m = nn.Sequential(nn.Linear(3, 4), nn.Tanh(), nn.Linear(4, 2))
        return m, zc.optim.AdamW(m.parameters(), lr=0.05, weight_decay=0.01)

    def train(m, opt, n):
        for _ in range(n):
            loss = F.cross_entropy(m(zc.Tensor(X)), y)
            opt.zero_grad()
            loss.backward()
            opt.step()

    a, opt_a = make()
    train(a, opt_a, 10)
    b, opt_b = make()
    train(b, opt_b, 5)
    model_state, opt_state = b.state_dict(), opt_b.state_dict()
    c, opt_c = make()
    c.load_state_dict(model_state)
    opt_c.load_state_dict(opt_state)
    train(c, opt_c, 5)
    assert all(np.array_equal(p, q) for p, q in zip(a.state_dict().values(), c.state_dict().values()))


# ---------------- 端到端学习 ----------------
def test_learns_xor():
    np.random.seed(0)
    X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
    y = np.array([0, 1, 1, 0])
    model = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 2))
    opt = zc.optim.AdamW(model.parameters(), lr=0.05)
    for _ in range(300):
        loss = F.cross_entropy(model(zc.Tensor(X)), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    assert (model(zc.Tensor(X)).data.argmax(1) == y).all(), f"XOR 没学会, loss={loss.item()}"


def test_attention_memorizes_sequence():
    """一层注意力 + 嵌入，记住一个短序列的下一个字符。"""
    np.random.seed(0)
    seq = np.array([0, 1, 2, 3, 2, 1, 0, 3])
    x, y = seq[None, :-1], seq[1:]

    class Tiny(nn.Module):
        def __init__(self):
            self.tok = nn.Embedding(4, 16)
            self.pos = nn.Embedding(8, 16)
            self.attn = nn.CausalSelfAttention(16, 2)
            self.ln = nn.LayerNorm(16)
            self.head = nn.Linear(16, 4)

        def forward(self, idx):
            T = idx.shape[1]
            h = self.tok(idx) + self.pos(np.arange(T))
            h = h + self.attn(self.ln(h))
            return self.head(h)

    model = Tiny()
    opt = zc.optim.AdamW(model.parameters(), lr=0.02)
    for _ in range(200):
        loss = F.cross_entropy(model(x).reshape(-1, 4), y)
        opt.zero_grad()
        loss.backward()
        zc.optim.clip_grad_norm(model.parameters(), 1.0)
        opt.step()
    assert loss.item() < 0.05, f"loss={loss.item()}"


if __name__ == "__main__":
    tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")]
    failed = 0
    for name, fn in tests:
        try:
            fn()
            print(f"  ok   {name}")
        except Exception as e:  # noqa: BLE001
            failed += 1
            print(f"  FAIL {name}: {type(e).__name__}: {e}")
    print(f"\n{len(tests) - failed}/{len(tests)} 通过")
    sys.exit(1 if failed else 0)