| |
| """J and R bi-encoder agreement. Does not recompute matmul or Goldilocks.""" |
| from __future__ import annotations |
|
|
| import json |
| import subprocess |
| import sys |
| import unittest |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| RESULT = ROOT / "results" / "biencoder_result.json" |
|
|
|
|
| class TestJAndRBiencoder(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| proc = subprocess.run( |
| [sys.executable, str(ROOT / "scripts" / "run_biencoder.py")], |
| cwd=str(ROOT), |
| capture_output=True, |
| text=True, |
| ) |
| cls.proc = proc |
| cls.data = json.loads(RESULT.read_text(encoding="utf-8")) |
|
|
| def test_runner_exits_zero(self): |
| self.assertEqual(self.proc.returncode, 0, self.proc.stdout + self.proc.stderr) |
| self.assertTrue(self.data["overall_pass"]) |
| self.assertEqual(self.data["correctness"]["failed"], 0) |
|
|
| def test_exact_head_matches_reference(self): |
| j = self.data["J"]["matrices"] |
| r = self.data["R"]["matrices"] |
| self.assertEqual(j["EXACT_HEAD"], j["REFERENCE"]) |
| self.assertEqual(r["EXACT_HEAD"], r["REFERENCE"]) |
| self.assertEqual(j["EXACT_HEAD"], r["EXACT_HEAD"]) |
| self.assertTrue(self.data["exact_head_match"]) |
| self.assertEqual(j["HAND_EXACT_HEAD"], j["HAND_REFERENCE"]) |
| self.assertEqual(j["HAND_EXACT_HEAD"], [[7, 7], [6, 9]]) |
| self.assertEqual(r["HAND_EXACT_HEAD"], [[7, 7], [6, 9]]) |
|
|
| def test_routed_error_is_recorded_not_required_zero(self): |
| routed = self.data["routed"] |
| self.assertIn("max_abs_error", routed) |
| self.assertIsInstance(routed["max_abs_error"], int) |
| self.assertFalse(routed["required_zero"]) |
| self.assertFalse(routed["gates_exit"]) |
| self.assertEqual(self.data["J"]["max_abs_error"], self.data["R"]["max_abs_error"]) |
| self.assertEqual(routed["max_abs_error"], self.data["J"]["max_abs_error"]) |
| self.assertEqual(self.data["J"]["matrices"]["PREDICTED"], self.data["R"]["matrices"]["PREDICTED"]) |
| self.assertNotEqual(self.data["J"]["matrices"]["PREDICTED"], self.data["J"]["matrices"]["EXACT_HEAD"]) |
| self.assertEqual(routed["hand_max_abs_error"], 2) |
| self.assertEqual(self.data["J"]["matrices"]["HAND_PREDICTED"], [[7, 5], [4, 9]]) |
| self.assertEqual(self.data["R"]["matrices"]["HAND_PREDICTED"], [[7, 5], [4, 9]]) |
|
|
| def test_permutation_inverse(self): |
| self.assertTrue(self.data["permutation_inverse"]) |
| self.assertTrue(self.data["J"]["permutation_inverse"]) |
| self.assertTrue(self.data["R"]["permutation_inverse"]) |
| self.assertEqual(self.data["J"]["matrices"]["WQ"], self.data["R"]["matrices"]["WK"]) |
|
|
| def test_goldilocks_small_integer_agreement(self): |
| self.assertTrue(self.data["goldilocks_small_agrees"]) |
| self.assertEqual(self.data["goldilocks_p"], "18446744069414584321") |
| j = self.data["J"]["matrices"] |
| r = self.data["R"]["matrices"] |
| self.assertEqual(j["FIELD_SMALL_PREDICTED"], j["PREDICTED"]) |
| self.assertEqual(j["FIELD_SMALL_EXACT_HEAD"], j["EXACT_HEAD"]) |
| self.assertEqual(r["FIELD_SMALL_PREDICTED"], r["PREDICTED"]) |
| self.assertEqual(r["FIELD_LARGE_EXACT_HEAD"], r["FIELD_LARGE_REFERENCE"]) |
| self.assertEqual(j["FIELD_LARGE_EXACT_HEAD"], r["FIELD_LARGE_EXACT_HEAD"]) |
| self.assertEqual(j["FIELD_LARGE_PREDICTED"], r["FIELD_LARGE_PREDICTED"]) |
| self.assertTrue(self.data["J"]["field_large_exact_match"]) |
| self.assertTrue(self.data["R"]["field_small_agrees"]) |
|
|
| def test_not_softmax_and_shapes(self): |
| self.assertEqual(self.data["shapes"], {"M": 4, "K": 4, "N": 4}) |
| self.assertIn("not softmax", self.data["note"]) |
| self.assertNotIn("trained softmax", self.data["note"]) |
|
|
| def test_sha512_dag_seal_agrees_and_mutation_fails(self): |
| seal = self.data["seal"] |
| self.assertTrue(seal["agree"]) |
| self.assertEqual(seal["J"], seal["R"]) |
| self.assertEqual(len(seal["J"]), 128) |
| self.assertEqual(seal["dag_order"], ["exact_head", "predicted", "weights", "prime"]) |
| self.assertTrue(seal["verify"]) |
| self.assertEqual(seal["verify_mutated"], 0) |
| self.assertEqual(seal["mutated_seal"]["J"], seal["mutated_seal"]["R"]) |
| self.assertNotEqual(seal["J"], seal["mutated_seal"]["J"]) |
| self.assertEqual(seal["neg_seal"]["J"], seal["neg_seal"]["R"]) |
| kat = self.data["sha512_kat"] |
| self.assertEqual( |
| kat["empty"], |
| "cf83e1357eefb8bdf1542850d66d8007d620e4050b5715dc83f4a921d36ce9ce47d0d13c5d85f2b0ff8318d2877eec2f63b931bd47417a81a538327af927da3e", |
| ) |
| self.assertEqual( |
| kat["abc"], |
| "ddaf35a193617abacc417349ae20413112e6fa4e89a97ea20a9eeee64b55d39a2192992a274fc1a836ba3c23a3feebbd454d4423643ce80e2a9ac94fa54ca49f", |
| ) |
| self.assertNotIn("sha256", self.data["note"].lower()) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
