File size: 1,709 Bytes
aeab535 |
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 49 |
import torch
from safetensors.torch import load_file
def load_model(path='model.safetensors'):
return load_file(path)
def priority_encode(bits, weights):
"""8-to-3 priority encoder with valid bit.
Returns (y2, y1, y0, valid) where y2y1y0 is the index of highest active input.
Highest index (x7) has highest priority.
"""
inp = torch.tensor([float(b) for b in bits])
# Compute winner for each input position
winners = []
for i in range(8):
win = int((inp * weights[f'winner{i}.weight']).sum() + weights[f'winner{i}.bias'] >= 0)
winners.append(win)
# Compute output bits from winners
win_137 = torch.tensor([float(winners[i]) for i in [1,3,5,7]])
y0 = int((win_137 * weights['y0.weight']).sum() + weights['y0.bias'] >= 0)
win_2367 = torch.tensor([float(winners[i]) for i in [2,3,6,7]])
y1 = int((win_2367 * weights['y1.weight']).sum() + weights['y1.bias'] >= 0)
win_4567 = torch.tensor([float(winners[i]) for i in [4,5,6,7]])
y2 = int((win_4567 * weights['y2.weight']).sum() + weights['y2.bias'] >= 0)
valid = int((inp * weights['valid.weight']).sum() + weights['valid.bias'] >= 0)
return y2, y1, y0, valid
if __name__ == '__main__':
w = load_model()
print('8-to-3 Priority Encoder (highest index wins)')
print('Input -> Index, Valid')
# Single-bit tests
for i in range(8):
bits = [0]*8
bits[i] = 1
y2, y1, y0, valid = priority_encode(bits, w)
print(f'x{i} only: {y2*4 + y1*2 + y0}, {valid}')
# No input
y2, y1, y0, valid = priority_encode([0]*8, w)
print(f'None: {y2*4 + y1*2 + y0}, valid={valid}')
|