File size: 1,672 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 46 47 48 49 50 51 52 | 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()
|