phanerozoic's picture
Upload folder using huggingface_hub
891df06 verified
import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def prefix_and(x3, x2, x1, x0, w):
"""4-bit prefix AND: y_i = AND(x3, x2, ..., x_i)"""
inp = torch.tensor([float(x3), float(x2), float(x1), float(x0)])
y3 = int((inp @ w['y3.weight'].T + w['y3.bias'] >= 0).item())
y2 = int((inp @ w['y2.weight'].T + w['y2.bias'] >= 0).item())
y1 = int((inp @ w['y1.weight'].T + w['y1.bias'] >= 0).item())
y0 = int((inp @ w['y0.weight'].T + w['y0.bias'] >= 0).item())
return y3, y2, y1, y0
if __name__ == '__main__':
w = load_model()
print('Prefix-AND Truth Table:')
print('x3x2x1x0 | y3y2y1y0 | meaning')
print('---------+----------+--------')
for i in [0b1111, 0b1110, 0b1101, 0b1011, 0b0111, 0b1100, 0b0000]:
x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
y3, y2, y1, y0 = prefix_and(x3, x2, x1, x0, w)
meaning = "all 1s" if y0 == 1 else f"first 0 at {3 - [y3,y2,y1,y0].index(0) if 0 in [y3,y2,y1,y0] else 'none'}"
print(f' {x3} {x2} {x1} {x0} | {y3} {y2} {y1} {y0} | {meaning}')