Add tests/test_loss_reweight.py: batch determinism + weighting semantics (runs in container)

1 files changed, 53 insertions, 0 deletions

diff --git a/tests/test_loss_reweight.py b/tests/test_loss_reweight.py
new file mode 100644
index 0000000..7df66c9
--- /dev/null
+++ b/tests/test_loss_reweight.py
@@ -0,0 +1,53 @@
+"""Minimal tests for the loss-reweighting experiment (no pytest needed).
+
+Run inside the meru container where torch exists:
+    python3 tests/test_loss_reweight.py
+Skips gracefully when torch is unavailable (e.g. on voidlaptop bare metal).
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+try:
+    import numpy as np
+    import torch
+    import torch.nn.functional as F
+    import loss_reweight as lr
+except ImportError:
+    print("SKIP: torch not available (run inside the meru container)")
+    sys.exit(0)
+
+DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def test_batch_determinism():
+    """Paired-model guarantee: identical minibatch order per seed, different across seeds."""
+    data = np.random.RandomState(0).randint(0, 65, 50000).astype(np.uint16)
+    g1a = torch.Generator().manual_seed(20260731)
+    g1b = torch.Generator().manual_seed(20260731)
+    x1, y1 = lr._batch(data, 128, 16, g1a, DEV)
+    x1b, y1b = lr._batch(data, 128, 16, g1b, DEV)
+    assert torch.equal(x1, x1b) and torch.equal(y1, y1b), "same seed must give same batch"
+    g2 = torch.Generator().manual_seed(20260732)
+    x2, _ = lr._batch(data, 128, 16, g2, DEV)
+    assert not torch.equal(x1, x2), "different seeds must give different batches"
+
+
+def test_weighting_semantics():
+    """control == plain CE; ctrl_random picks a fixed set of non-q positions."""
+    q_id = 3
+    y = torch.randint(0, 65, (2, 4), device=DEV)
+    logits = torch.randn(2, 4, 65, device=DEV)
+    l_ctl = lr._weighted_loss(logits, y, 'control', q_id, 0, 65, DEV)
+    l_ce = F.cross_entropy(logits.view(-1, 65), y.view(-1))
+    assert abs(l_ctl.item() - l_ce.item()) < 1e-6, "control must equal plain CE"
+    l_r1 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV)
+    l_r2 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV)
+    assert abs(l_r1.item() - l_r2.item()) < 1e-9, "ctrl_random must be deterministic"
+
+
+if __name__ == '__main__':
+    test_batch_determinism()
+    test_weighting_semantics()
+    print(f"PASS ({DEV})")