File size: 1,083 Bytes
26e7af6 | 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 | import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def ffs8(bits, w):
"""Find first set bit (1-indexed). Returns 0 if no bits set."""
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)
f0 = 1 if (p[0] or p[2] or p[4] or p[6]) else 0
f1 = 1 if (p[1] or p[2] or p[5] or p[6]) else 0
f2 = 1 if (p[3] or p[4] or p[5] or p[6]) else 0
f3 = p[7]
return f3, f2, f1, f0
if __name__ == '__main__':
w = load_model()
print('FFS8 selected tests:')
print('Input | FFS | f3f2f1f0')
print('---------+-----+---------')
test_vals = [0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x06, 0xF0, 0xFF]
for val in test_vals:
bits = [(val >> j) & 1 for j in range(8)]
f3, f2, f1, f0 = ffs8(bits, w)
pos = 8*f3 + 4*f2 + 2*f1 + f0
print(f'{val:08b} | {pos} | {f3}{f2}{f1}{f0}')
|