File size: 1,277 Bytes
324ecd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | import unittest
import torch
from beamboot.cpu import CompiledCPU, Instruction, HALT, VOCAB, encode, output_bytes, print_program, run, transition
class CPUTests(unittest.TestCase):
def test_all_states_match_reference(self):
program = [Instruction('LDI', 3), Instruction('SUB', 1), Instruction('JZ', 4), Instruction('JMP', 1), Instruction('OUT'), Instruction('HALT')]
model = CompiledCPU(program)
with torch.inference_mode():
result = model(torch.arange(VOCAB)).argmax(dim=-1).tolist()
self.assertEqual(result, [transition(program, i) for i in range(VOCAB)])
def test_boot_text(self):
model = CompiledCPU(print_program('OK\n'))
self.assertEqual(output_bytes(list(run(model))), b'OK\n')
def test_loop_and_output(self):
program = [Instruction('LDI', 3), Instruction('OUT'), Instruction('SUB', 1), Instruction('JZ', 5), Instruction('JMP', 1), Instruction('HALT')]
result = list(run(CompiledCPU(program)))
self.assertEqual(result[-1], HALT)
self.assertEqual([t // 2 % 16 for t in result if t != HALT and t % 2], [3, 2, 1])
def test_infinite_loop_budget(self):
with self.assertRaises(RuntimeError):
list(run(CompiledCPU([Instruction('JMP', 0)]), 3))
|