File size: 1,149 Bytes
905a4f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def ctz8(bits, w):
"""Count trailing zeros in 8-bit input. bits[0] is LSB."""
inp = torch.tensor([float(b) for b in bits])
p = []
for i in range(8):
pi = int((inp @ w[f'p{i}.weight'].T + w[f'p{i}.bias'] >= 0).item())
p.append(pi)
pZ = int((inp @ w['pZ.weight'].T + w['pZ.bias'] >= 0).item())
c0 = 1 if (p[1] or p[3] or p[5] or p[7]) else 0
c1 = 1 if (p[2] or p[3] or p[6] or p[7]) else 0
c2 = 1 if (p[4] or p[5] or p[6] or p[7]) else 0
c3 = pZ
return c3, c2, c1, c0
if __name__ == '__main__':
w = load_model()
print('CTZ8 selected tests:')
print('Input | CTZ | c3c2c1c0')
print('---------+-----+---------')
test_vals = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x00, 0x06, 0xF0, 0xFF]
for val in test_vals:
bits = [(val >> j) & 1 for j in range(8)]
c3, c2, c1, c0 = ctz8(bits, w)
count = 8*c3 + 4*c2 + 2*c1 + c0
print(f'{val:08b} | {count} | {c3}{c2}{c1}{c0}')
|