File size: 1,225 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 | import random
import unittest
import torch
from beamboot.tensor_memory import TensorMemory
class TensorMemoryTests(unittest.TestCase):
def test_updates_match_reference(self):
rng = random.Random(812)
mem = TensorMemory(max_writes=256)
reference = [0] * 16
for _ in range(256):
a, v = rng.randrange(16), rng.randrange(256)
mem.write(a, v)
reference[a] = v
self.assertEqual(mem.read_byte(a), v)
for a in range(16):
self.assertEqual(mem.read_byte(a), reference[a])
def test_window_loses_old_memory(self):
mem = TensorMemory(window=2)
mem.write(1, 42)
mem.write(2, 20)
mem.write(3, 30)
self.assertNotEqual(mem.read_byte(1), 42)
def test_write_limit_and_validation(self):
mem = TensorMemory(max_writes=1)
with self.assertRaises(ValueError): mem.write(16, 1)
mem.write(0, 255)
with self.assertRaises(ValueError): mem.write(0, 2)
def test_bfloat16_is_measured_not_assumed(self):
mem = TensorMemory(max_writes=4, dtype=torch.bfloat16)
mem.write(1, 11)
mem.write(1, 22)
self.assertEqual(mem.read_byte(1), 22)
|