| import random |
| import unittest |
| from beamboot.memory import AttentionMemory |
|
|
|
|
| class TestAttentionMemory(unittest.TestCase): |
| def test_latest_write_and_missing(self): |
| mem = AttentionMemory() |
| self.assertIsNone(mem.read(3)) |
| mem.write(3, 12) |
| mem.write(5, 99) |
| mem.write(3, 231) |
| self.assertEqual(mem.read_byte(3), 231) |
| self.assertEqual(mem.read_byte(5), 99) |
| self.assertIsNone(mem.read(2)) |
|
|
| def test_against_reference(self): |
| randomizer = random.Random(42) |
| mem, reference = AttentionMemory(temperature=0.00001), {} |
| for _ in range(500): |
| address = randomizer.randrange(16) |
| if randomizer.random() < .55: |
| value = randomizer.randrange(256) |
| reference[address] = value |
| mem.write(address, value) |
| else: |
| self.assertEqual(mem.read_byte(address), reference.get(address)) |
|
|
| def test_sliding_window_forgets(self): |
| mem = AttentionMemory(window=2) |
| mem.write(1, 42) |
| mem.write(2, 10) |
| mem.write(3, 20) |
| self.assertIsNone(mem.read(1)) |
| self.assertEqual(mem.read_byte(2), 10) |
|
|
| def test_softmax_is_not_exact(self): |
| mem = AttentionMemory(temperature=1.0) |
| mem.write(1, 0) |
| mem.write(1, 255) |
| self.assertNotEqual(mem.read_byte(1), 255) |
|
|
| def test_validation(self): |
| mem = AttentionMemory() |
| with self.assertRaises(ValueError): mem.write(16, 1) |
| with self.assertRaises(ValueError): mem.write(1, 256) |
| with self.assertRaises(ValueError): mem.read(-1) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
