h3-blackwell-runtime/research/vortex_exact_attention/tests/test_models.py
2026-08-26 13:39:20 +07:00

48 lines
1.9 KiB
Python

import importlib.util
import copy
from pathlib import Path
import unittest
class ModelContracts(unittest.TestCase):
@classmethod
def setUpClass(cls):
path = Path(__file__).resolve().parents[1] / "tools/validate_models.py"
spec = importlib.util.spec_from_file_location("validate_vortex_models", path)
cls.module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(cls.module)
def test_retained_reference_and_architecture_models(self):
self.module.validate()
def test_rejects_same_warp_state_ownership(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][0]["handoff"]["same_warp"] = True
with self.assertRaisesRegex(RuntimeError, "one warp"):
self.module.validate_documents(inventory, changed)
def test_rejects_inconsistent_shared_memory(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][1]["shared_kib"]["total"] = [52, 56]
with self.assertRaisesRegex(RuntimeError, "VEA-B shared-memory"):
self.module.validate_documents(inventory, changed)
def test_rejects_inconsistent_screening_arithmetic(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["candidates"][0]["mainloop_ms_projected"] = [192, 220]
with self.assertRaisesRegex(RuntimeError, "screening arithmetic"):
self.module.validate_documents(inventory, changed)
def test_rejects_measured_result_label(self):
inventory, model = self.module.load_models()
changed = copy.deepcopy(model)
changed["status"] = "achieved"
with self.assertRaisesRegex(RuntimeError, "measurements"):
self.module.validate_documents(inventory, changed)
if __name__ == "__main__":
unittest.main()