File size: 1,144 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 | import random
import unittest
import torch
from beamboot.machine import NeuralMachine, State, decode_output, reference_step, run, state_bits
class MachineTests(unittest.TestCase):
def test_all_states(self):
program = [('LDI', 7), ('STM', 0), ('SUB', 1), ('JZ', 6), ('ADD', 2), ('JMP', 2), ('LDM', 0), ('OUT', 0), ('HALT', 0)]
model = NeuralMachine(program)
states = [State(pc, acc, ram, out) for pc in range(16) for acc in range(16) for ram in range(16) for out in (False, True)]
with torch.inference_mode():
for start in range(0, len(states), 256):
chunk = states[start:start+256]
actual = model(torch.stack([state_bits(s) for s in chunk]))
self.assertEqual([decode_output(row) for row in actual], [reference_step(program, s) for s in chunk])
def test_writable_memory(self):
program = [('LDI', 11), ('STM', 0), ('LDI', 0), ('LDM', 0), ('OUT', 0), ('HALT', 0)]
states = list(run(NeuralMachine(program)))
self.assertEqual([(s.acc, s.ram) for s in states if s.out], [(11, 11)])
self.assertTrue(states[-1].halt)
|