| 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) |
|
|
