File size: 2,640 Bytes
1529ce2 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import torch
from safetensors.torch import save_file
weights = {}
# 2x2 Array Multiplier
# Uses regular array structure of AND gates and adders
# Inputs: A1,A0, B1,B0 (4 inputs)
# Outputs: P3,P2,P1,P0 (4 outputs)
def add_and(name, idx_a, idx_b, n_inputs):
w = [0.0] * n_inputs
w[idx_a] = 1.0
w[idx_b] = 1.0
weights[f'{name}.weight'] = torch.tensor([w], dtype=torch.float32)
weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)
def add_xor_direct(name):
weights[f'{name}.or.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights[f'{name}.or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights[f'{name}.nand.weight'] = torch.tensor([[-1.0, -1.0]], dtype=torch.float32)
weights[f'{name}.nand.bias'] = torch.tensor([1.0], dtype=torch.float32)
weights[f'{name}.and.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights[f'{name}.and.bias'] = torch.tensor([-2.0], dtype=torch.float32)
def add_ha_carry(name):
weights[f'{name}.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)
# Row 0: A[1:0] * B0
add_and('r0c0', 1, 3, 4) # A0*B0 -> P0
add_and('r0c1', 0, 3, 4) # A1*B0
# Row 1: A[1:0] * B1, shifted and added
add_and('r1c0', 1, 2, 4) # A0*B1
add_and('r1c1', 0, 2, 4) # A1*B1
# Array cell (0,1): r0c1 + r1c0 -> half adder
add_xor_direct('cell01_sum')
add_ha_carry('cell01_carry')
# Array cell (1,1): r1c1 + carry -> half adder
add_xor_direct('cell11_sum')
add_ha_carry('cell11_carry')
save_file(weights, 'model.safetensors')
def array_mult(a1, a0, b1, b0):
# Partial products
r0c0 = a0 & b0 # P0
r0c1 = a1 & b0
r1c0 = a0 & b1
r1c1 = a1 & b1
# Array addition
p0 = r0c0
p1 = r0c1 ^ r1c0
c1 = r0c1 & r1c0
p2 = r1c1 ^ c1
p3 = r1c1 & c1
return p3, p2, p1, p0
print("Verifying 2x2 array multiplier...")
errors = 0
for a in range(4):
for b in range(4):
a1, a0 = (a >> 1) & 1, a & 1
b1, b0 = (b >> 1) & 1, b & 1
p3, p2, p1, p0 = array_mult(a1, a0, b1, b0)
result = p3*8 + p2*4 + p1*2 + p0
expected = a * b
if result != expected:
errors += 1
print(f"ERROR: {a}*{b} = {result}, expected {expected}")
if errors == 0:
print("All 16 test cases passed!")
else:
print(f"FAILED: {errors} errors")
mag = sum(t.abs().sum().item() for t in weights.values())
print(f"Magnitude: {mag:.0f}")
print(f"Parameters: {sum(t.numel() for t in weights.values())}")
|