|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import os |
| import json |
| import safetensors.torch |
|
|
|
|
| class AdderLayer(nn.Module): |
| def __init__(self, i_dim: int, o_dim: int): |
| super().__init__() |
| self.linear = nn.Linear(i_dim, o_dim) |
| |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.linear(x) |
| x = F.relu(x) |
| x_shape = x.shape |
| x = x.reshape(x_shape[:-1] + (-1, 2)) |
| x = x / x.norm(dim=-1, keepdim=True) |
| x = x.reshape(x_shape) |
| return x |
|
|
|
|
| class AdderEncoder(nn.Module): |
| def __init__(self, bits: int = 32): |
| super().__init__() |
| self.bits = bits |
|
|
| def forward(self, x: torch.LongTensor) -> torch.Tensor: |
| o = torch.zeros(x.shape + (2 * self.bits,), device=x.device) |
| for i in range(self.bits): |
| v = (x >> i) & 1 |
| o[..., i * 2] = 1 - v |
| o[..., i * 2 + 1] = v |
| return o |
|
|
| def extra_repr(self) -> str: |
| return f'bits={self.bits}' |
|
|
|
|
| class AdderDecoder(nn.Module): |
| def __init__(self, bits: int = 32): |
| super().__init__() |
| self.bits = bits |
| |
| def forward(self, x: torch.Tensor) -> torch.LongTensor: |
| o = torch.zeros(x.shape[:-1], device=x.device, dtype=torch.long) |
| for i in range(self.bits): |
| v = x[..., i * 2 + 1] > x[..., i * 2] |
| o = o | (v << i) |
| return o |
|
|
|
|
| class Adder(nn.Module): |
| def __init__(self, layer_dims: list[int], bits: int = 32): |
| super().__init__() |
| self.layer_dims = layer_dims |
| self.bits = bits |
| self.encoder = AdderEncoder(bits) |
| self.encoder_c = AdderEncoder(1) |
| self.decoder = AdderDecoder(bits) |
| self.decoder_c = AdderDecoder(1) |
| self.layers = nn.ModuleList() |
| for i_dim, o_dim in zip(layer_dims[:-1], layer_dims[1:]): |
| self.layers.append(AdderLayer(i_dim, o_dim)) |
|
|
| @property |
| def config(self) -> dict: |
| return { |
| 'layer_dims': self.layer_dims, |
| 'bits': self.bits |
| } |
|
|
| @classmethod |
| def from_pretrained(cls, filepath: str): |
| with open(os.path.join(filepath, 'config.json'), 'r') as f: |
| config = json.load(f) |
| model = cls(**config) |
| state_dict = safetensors.torch.load_file( |
| os.path.join(filepath, 'model.safetensors')) |
| model.load_state_dict(state_dict) |
| model.requires_grad_(False) |
| return model |
|
|
| def forward(self, a: torch.LongTensor, b: torch.LongTensor) -> torch.LongTensor: |
| assert (0 <= a < 2 ** self.bits).all() |
| assert (0 <= b < 2 ** self.bits).all() |
| a = self.encoder(a) |
| b = self.encoder(b) |
| c = self.encoder_c(torch.tensor(0, device=a.device)) |
| x = torch.cat([a, b, c], dim=-1) |
| for m in self.layers: |
| x = m(x) |
| x, c = x.split([2 * self.bits, 2], dim=-1) |
| x = self.decoder(x) |
| c = self.decoder_c(c) |
| if (c > 0).any(): |
| raise ValueError("Carry out is not 0") |
| return x |
|
|
|
|
| def main(): |
| model = Adder.from_pretrained(os.path.dirname(__file__)) |
| print(model(torch.tensor(3123), torch.tensor(5929))) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
