| import torch | |
| from safetensors.torch import load_file | |
| def load_model(path='model.safetensors'): | |
| return load_file(path) | |
| def xor_gate(a, b, w, idx): | |
| inp = torch.tensor([float(a), float(b)]) | |
| l1 = (inp @ w[f'xor{idx}.layer1.weight'].T + w[f'xor{idx}.layer1.bias'] >= 0).float() | |
| return int((l1 @ w[f'xor{idx}.layer2.weight'].T + w[f'xor{idx}.layer2.bias'] >= 0).item()) | |
| def equal(a, b, weights): | |
| """8-bit equality comparator. | |
| a, b: lists of 8 bits each (LSB first) | |
| Returns: 1 if a == b, 0 otherwise | |
| """ | |
| xors = [xor_gate(a[i], b[i], weights, i) for i in range(8)] | |
| xor_vec = torch.tensor([float(x) for x in xors]) | |
| return int((xor_vec @ weights['nor.weight'].T + weights['nor.bias'] >= 0).item()) | |
| if __name__ == '__main__': | |
| w = load_model() | |
| print('8-bit Equal Comparator') | |
| print('a == b tests:') | |
| tests = [(0, 0), (0, 1), (127, 127), (127, 128), (255, 255), (255, 0), (100, 100)] | |
| for a_val, b_val in tests: | |
| a = [(a_val >> i) & 1 for i in range(8)] | |
| b = [(b_val >> i) & 1 for i in range(8)] | |
| result = equal(a, b, w) | |
| print(f'{a_val:3d} == {b_val:3d} = {result}') | |