| 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)]) |
|
|
| |
| 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()) |
|
|
| |
| 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, 1, 0, 0, 0, 0, 0, 0), |
| (0, 0, 0, 0, 1, 0, 0, 0), |
| (0, 0, 0, 0, 0, 0, 0, 1), |
| (0, 0, 0, 0, 0, 0, 0, 0), |
| ] |
| 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}') |
|
|
