| """Bounded cache memory using PyTorch's scaled-dot-product attention. |
| |
| This is a constructed attention experiment, not Beam's attention implementation. |
| """ |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| class TensorMemory: |
| def __init__(self, address_bits=4, max_writes=256, window=None, dtype=torch.float32): |
| if address_bits < 1 or max_writes < 1 or (window is not None and window < 1): |
| raise ValueError('invalid configuration') |
| self.address_bits = address_bits |
| self.max_writes = max_writes |
| self.window = window |
| self.dtype = dtype |
| self.keys = [] |
| self.values = [] |
| |
| for address in range(2**address_bits): |
| self._append(address, 0, 0) |
| self.writes = 0 |
|
|
| def _bits(self, address): |
| if not isinstance(address, int) or not 0 <= address < 2**self.address_bits: |
| raise ValueError('address out of range') |
| return [1.0 if address & (1 << i) else -1.0 for i in range(self.address_bits)] |
|
|
| def _append(self, address, value, position): |
| |
| |
| self.keys.append(self._bits(address) + [float(position)]) |
| self.values.append([float(value)]) |
|
|
| def write(self, address, value): |
| if self.writes >= self.max_writes: |
| raise ValueError('write budget exhausted') |
| if not isinstance(value, int) or not 0 <= value <= 255: |
| raise ValueError('value must be a byte') |
| self._bits(address) |
| self.writes += 1 |
| self._append(address, value, self.writes) |
|
|
| def read(self, address): |
| bits = self._bits(address) |
| |
| |
| |
| query = bits + [1.0] |
| keys = self.keys[-self.window:] if self.window else self.keys |
| values = self.values[-self.window:] if self.window else self.values |
| q = torch.tensor(query, dtype=self.dtype).reshape(1, 1, 1, -1) |
| k = torch.tensor(keys, dtype=self.dtype).reshape(1, 1, len(keys), -1) |
| v = torch.tensor(values, dtype=self.dtype).reshape(1, 1, len(values), 1) |
| q[..., :-1] *= 16 * (self.max_writes + 1) |
| q[..., -1] *= 16 |
| with torch.no_grad(): |
| return float(F.scaled_dot_product_attention(q, k, v, scale=1.0).item()) |
|
|
| def read_byte(self, address): |
| return round(self.read(address)) |
|
|
