File size: 2,465 Bytes
cdd5bc3 |
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 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def clz8(a7, a6, a5, a4, a3, a2, a1, a0, weights):
"""8-bit count leading zeros. Returns 4-bit binary encoding of 0-8."""
inp = torch.tensor([float(a7), float(a6), float(a5), float(a4),
float(a3), float(a2), float(a1), float(a0)])
# Layer 1: priority detection
has7 = int((inp @ weights['has7.weight'].T + weights['has7.bias'] >= 0).item())
has6_first = int((inp @ weights['has6_first.weight'].T + weights['has6_first.bias'] >= 0).item())
has5_first = int((inp @ weights['has5_first.weight'].T + weights['has5_first.bias'] >= 0).item())
has4_first = int((inp @ weights['has4_first.weight'].T + weights['has4_first.bias'] >= 0).item())
has3_first = int((inp @ weights['has3_first.weight'].T + weights['has3_first.bias'] >= 0).item())
has2_first = int((inp @ weights['has2_first.weight'].T + weights['has2_first.bias'] >= 0).item())
has1_first = int((inp @ weights['has1_first.weight'].T + weights['has1_first.bias'] >= 0).item())
has0_first = int((inp @ weights['has0_first.weight'].T + weights['has0_first.bias'] >= 0).item())
all_zero = int((inp @ weights['all_zero.weight'].T + weights['all_zero.bias'] >= 0).item())
# Layer 2: binary encoding
l1 = torch.tensor([float(has7), float(has6_first), float(has5_first), float(has4_first),
float(has3_first), float(has2_first), float(has1_first), float(has0_first),
float(all_zero)])
y0 = int((l1 @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
y1 = int((l1 @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
y2 = int((l1 @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
y3 = int((l1 @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())
return [y3, y2, y1, y0]
if __name__ == '__main__':
w = load_model()
print('clz8 examples:')
test_cases = [
(1, 0, 0, 0, 0, 0, 0, 0), # 0 leading zeros
(0, 1, 0, 0, 0, 0, 0, 0), # 1 leading zero
(0, 0, 0, 0, 1, 0, 0, 0), # 4 leading zeros
(0, 0, 0, 0, 0, 0, 0, 1), # 7 leading zeros
(0, 0, 0, 0, 0, 0, 0, 0), # 8 leading zeros
]
for bits in test_cases:
result = clz8(*bits, w)
clz_val = result[0]*8 + result[1]*4 + result[2]*2 + result[3]
print(f' {bits} -> {result} = {clz_val}')
|