1 files changed, 26 insertions, 0 deletions
diff --git a/tests/test_scaling.py b/tests/test_scaling.py
new file mode 100644
index 0000000..b66a05c
--- /dev/null
+++ b/tests/test_scaling.py
@@ -0,0 +1,26 @@
+import os, sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
+import torch
+import torch.nn as nn
+from zbp_scaling.model import ScalingLM
+from zbp_scaling.zbp import ZBPConfig, zbp_blocks
+
+
+def test_exact_mode_matches_bp():
+    torch.manual_seed(0)
+    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="bp"))
+    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
+    g_bp = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
+    for b in zbp_blocks(m):
+        b.cfg = ZBPConfig(mode="exact")
+    g = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
+    err = max((a - b).norm().item() / (b.norm().item() + 1e-12) for a, b in zip(g, g_bp))
+    assert err < 1e-5, err
+
+
+def test_zbp_cd_trains_shape():
+    torch.manual_seed(0)
+    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="cd", n_probes=4, eps=0.1, probe_chunk=4))
+    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
+    nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()).backward()
+    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters())