create_safetensors.py

import torch
from safetensors.torch import save_file

weights = {}

# BCD Adder (single digit)
# Inputs: a3,a2,a1,a0, b3,b2,b1,b0, cin (9 inputs)
# Outputs: s3,s2,s1,s0, cout (5 outputs)
# Valid inputs: 0-9 for A and B
#
# Algorithm:
# 1. Binary add: Z = A + B + Cin
# 2. If Z > 9 or carry: add 6 for correction
# 3. Decimal carry if Z > 9 or binary carry

def add_xor(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_maj(name):
    weights[f'{name}.weight'] = torch.tensor([[1.0, 1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_fa(name):
    add_xor(f'{name}.xor1')
    add_xor(f'{name}.sum')
    add_maj(f'{name}.carry')

def add_ha(name):
    add_xor(f'{name}.sum')
    weights[f'{name}.carry.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
    weights[f'{name}.carry.bias'] = torch.tensor([-2.0], dtype=torch.float32)

def add_and(name, n_inputs):
    weights[f'{name}.weight'] = torch.tensor([[1.0] * n_inputs], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-float(n_inputs)], dtype=torch.float32)

def add_or(name, n_inputs):
    weights[f'{name}.weight'] = torch.tensor([[1.0] * n_inputs], dtype=torch.float32)
    weights[f'{name}.bias'] = torch.tensor([-1.0], dtype=torch.float32)

# Stage 1: 4-bit binary adder
add_fa('bin_fa0')
add_fa('bin_fa1')
add_fa('bin_fa2')
add_fa('bin_fa3')

# Stage 2: Detect if > 9 or carry
# Z > 9 means: Z3·Z2 + Z3·Z1 = Z3·(Z2+Z1)
# Or equivalently: (Z >= 10) = Z3·Z2 + Z3·Z1
add_and('detect_z3z2', 2)
add_and('detect_z3z1', 2)
add_or('detect_gt9', 2)
add_or('need_correct', 2)

# Stage 3: Conditional add 6 (only bits 1,2 affected, bit 0 unchanged)
# If correction needed: add 0110 to Z
add_ha('corr_ha1')
add_ha('corr_ha2')
add_xor('corr_xor3')

save_file(weights, 'model.safetensors')

def eval_xor(a, b):
    or_out = int(a + b >= 1)
    nand_out = int(-a - b + 1 >= 0)
    return int(or_out + nand_out >= 2)

def eval_maj(a, b, c):
    return int(a + b + c >= 2)

def eval_fa(a, b, cin):
    x1 = eval_xor(a, b)
    s = eval_xor(x1, cin)
    c = eval_maj(a, b, cin)
    return s, c

def eval_ha(a, b):
    return eval_xor(a, b), int(a + b >= 2)

def bcd_add(a3, a2, a1, a0, b3, b2, b1, b0, cin):
    z0, c0 = eval_fa(a0, b0, cin)
    z1, c1 = eval_fa(a1, b1, c0)
    z2, c2 = eval_fa(a2, b2, c1)
    z3, c3 = eval_fa(a3, b3, c2)
    gt9_a = z3 and z2
    gt9_b = z3 and z1
    gt9 = gt9_a or gt9_b
    need_corr = gt9 or c3
    if need_corr:
        s0 = z0
        s1, c_s1 = eval_ha(z1, 1)
        s2, c_s2 = eval_ha(z2, c_s1)
        s2, c_s2b = eval_ha(s2, 1)
        s3 = eval_xor(z3, c_s2 or c_s2b)
        cout = 1
    else:
        s0, s1, s2, s3 = z0, z1, z2, z3
        cout = 0
    return s0, s1, s2, s3, cout

print("Verifying BCD Adder...")
errors = 0
for a in range(10):
    for b in range(10):
        for cin in range(2):
            a3, a2, a1, a0 = (a>>3)&1, (a>>2)&1, (a>>1)&1, a&1
            b3, b2, b1, b0 = (b>>3)&1, (b>>2)&1, (b>>1)&1, b&1
            s0, s1, s2, s3, cout = bcd_add(a3, a2, a1, a0, b3, b2, b1, b0, cin)
            result = s0 + (s1 << 1) + (s2 << 2) + (s3 << 3)
            decimal_result = result + cout * 10
            expected = a + b + cin
            if decimal_result != expected:
                errors += 1
                if errors <= 5:
                    print(f"ERROR: {a}+{b}+{cin} = {decimal_result} (s={result},c={cout}), expected {expected}")

if errors == 0:
    print("All 200 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())}")