File size: 2,279 Bytes
5334ac8 | 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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | import torch
from safetensors.torch import save_file
weights = {}
# 4-bit Buffer (identity function)
# Inputs: x3, x2, x1, x0
# Outputs: y3, y2, y1, y0 (same as inputs)
#
# Each output is a threshold neuron that fires when input >= 1
# y_i = 1 iff x_i = 1
# y0 = x0
weights['y0.weight'] = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=torch.float32)
weights['y0.bias'] = torch.tensor([-1.0], dtype=torch.float32)
# y1 = x1
weights['y1.weight'] = torch.tensor([[0.0, 0.0, 1.0, 0.0]], dtype=torch.float32)
weights['y1.bias'] = torch.tensor([-1.0], dtype=torch.float32)
# y2 = x2
weights['y2.weight'] = torch.tensor([[0.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['y2.bias'] = torch.tensor([-1.0], dtype=torch.float32)
# y3 = x3
weights['y3.weight'] = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['y3.bias'] = torch.tensor([-1.0], dtype=torch.float32)
save_file(weights, 'model.safetensors')
def buffer4(x3, x2, x1, x0):
inp = torch.tensor([float(x3), float(x2), float(x1), float(x0)])
y0 = int((inp @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
y1 = int((inp @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
y2 = int((inp @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
y3 = int((inp @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())
return y3, y2, y1, y0
print("Verifying 4-bit Buffer...")
errors = 0
for i in range(16):
x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
y3, y2, y1, y0 = buffer4(x3, x2, x1, x0)
if (y3, y2, y1, y0) != (x3, x2, x1, x0):
errors += 1
print(f"ERROR: ({x3},{x2},{x1},{x0}) -> ({y3},{y2},{y1},{y0})")
if errors == 0:
print("All 16 test cases passed!")
else:
print(f"FAILED: {errors} errors")
print("\nTruth Table:")
print("x3 x2 x1 x0 | y3 y2 y1 y0")
print("-" * 26)
for i in range(16):
x3, x2, x1, x0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
y3, y2, y1, y0 = buffer4(x3, x2, x1, x0)
print(f" {x3} {x2} {x1} {x0} | {y3} {y2} {y1} {y0}")
mag = sum(t.abs().sum().item() for t in weights.values())
print(f"\nMagnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")
print(f"Neurons: {len([k for k in weights.keys() if 'weight' in k])}")
|