File size: 2,075 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import unittest
import torch
from beamboot.sequence_cpu import SequenceCPU, SequenceRuntime, event
from beamboot.monitor import Monitor, firmware
class SequenceTests(unittest.TestCase):
def test_full_attention_is_causal_and_last_decode_matches(self):
model = SequenceCPU([('LOAD', 2), ('HALT', 0)])
events = [event(address=i, valid=True) for i in range(16)]
events += [event(address=2, value=7, valid=True, position=16), event(position=17)]
x = torch.tensor([events], dtype=torch.float32)
with torch.inference_mode():
first = model(x)
cached = model(x, last_only=True)
changed = x.clone()
changed[0, -1, 14:18] = 1
changed_result = model(changed)
self.assertTrue(torch.equal(first[:, -1:], cached))
self.assertTrue(torch.equal(first[:, :-1], changed_result[:, :-1]))
self.assertEqual(first[0, -1, 5:9].tolist(), [1, 1, 1, 0])
def test_stores_and_reloads(self):
program = [('LDI', 9), ('STORE', 4), ('LDI', 2), ('STORE', 5),
('LDI', 0), ('LOAD', 4), ('OUT', 0), ('LOAD', 5), ('OUT', 0), ('HALT', 0)]
runtime = SequenceRuntime(SequenceCPU(program))
steps = []
while not runtime.halted: steps.append(runtime.step())
self.assertEqual([s.acc for s in steps if s.output], [9, 2])
self.assertEqual([(s.address, s.value) for s in steps if s.write], [(4, 9), (5, 2)])
def test_monitor_firmware(self):
monitor = Monitor()
monitor.runtime = SequenceRuntime(SequenceCPU(firmware()))
self.assertEqual(monitor.boot(), b'>')
self.assertEqual(monitor.command(1), b'OK')
self.assertEqual(monitor.command(2), b'HI')
self.assertEqual(monitor.command(3), b'')
self.assertEqual(monitor.command(1), b'OK')
def test_bounded_context(self):
machine = SequenceRuntime(SequenceCPU([('JMP', 0)], max_events=18))
machine.step(); machine.step()
with self.assertRaisesRegex(RuntimeError, 'budget'): machine.step()
|