create_safetensors.py

import torch
from safetensors.torch import save_file

weights = {}

# Input: [a3, a2, a1, a0]
# Output: input + 1 (mod 16)
# y0 = NOT(a0)
# y1 = a1 XOR a0
# y2 = a2 XOR (a1 AND a0)
# y3 = a3 XOR (a2 AND a1 AND a0)

# Layer 1: NOT, carries, and XOR components
# y0 = NOT(a0)
weights['y0.weight'] = torch.tensor([[0.0, 0.0, 0.0, -1.0]], dtype=torch.float32)
weights['y0.bias'] = torch.tensor([0.0], dtype=torch.float32)

# c2 = a1 AND a0
weights['c2.weight'] = torch.tensor([[0.0, 0.0, 1.0, 1.0]], dtype=torch.float32)
weights['c2.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# c3 = a2 AND a1 AND a0
weights['c3.weight'] = torch.tensor([[0.0, 1.0, 1.0, 1.0]], dtype=torch.float32)
weights['c3.bias'] = torch.tensor([-3.0], dtype=torch.float32)

# XOR(a1, a0) components
weights['y1_or.weight'] = torch.tensor([[0.0, 0.0, 1.0, 1.0]], dtype=torch.float32)
weights['y1_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['y1_nand.weight'] = torch.tensor([[0.0, 0.0, -1.0, -1.0]], dtype=torch.float32)
weights['y1_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)

# Layer 2: y1 AND, XOR(a2,c2) and XOR(a3,c3) components
# Input order for layer 2: [a3, a2, c2, c3, y1_or, y1_nand]
# y1 = AND(y1_or, y1_nand)
weights['y1.weight'] = torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 1.0]], dtype=torch.float32)
weights['y1.bias'] = torch.tensor([-2.0], dtype=torch.float32)

# y2 = XOR(a2, c2): OR(a2, c2) and NAND(a2, c2)
weights['y2_or.weight'] = torch.tensor([[0.0, 1.0, 1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['y2_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['y2_nand.weight'] = torch.tensor([[0.0, -1.0, -1.0, 0.0, 0.0, 0.0]], dtype=torch.float32)
weights['y2_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)

# y3 = XOR(a3, c3): OR(a3, c3) and NAND(a3, c3)
weights['y3_or.weight'] = torch.tensor([[1.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
weights['y3_or.bias'] = torch.tensor([-1.0], dtype=torch.float32)
weights['y3_nand.weight'] = torch.tensor([[-1.0, 0.0, 0.0, -1.0, 0.0, 0.0]], dtype=torch.float32)
weights['y3_nand.bias'] = torch.tensor([1.0], dtype=torch.float32)

# Layer 3: y2 and y3 final AND
weights['y2.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['y2.bias'] = torch.tensor([-2.0], dtype=torch.float32)
weights['y3.weight'] = torch.tensor([[1.0, 1.0]], dtype=torch.float32)
weights['y3.bias'] = torch.tensor([-2.0], dtype=torch.float32)

save_file(weights, 'model.safetensors')

# Verify
def incrementer4(a3, a2, a1, a0):
    inp = torch.tensor([float(a3), float(a2), float(a1), float(a0)])

    # Layer 1
    y0 = int((inp @ weights['y0.weight'].T + weights['y0.bias'] >= 0).item())
    c2 = int((inp @ weights['c2.weight'].T + weights['c2.bias'] >= 0).item())
    c3 = int((inp @ weights['c3.weight'].T + weights['c3.bias'] >= 0).item())
    y1_or = int((inp @ weights['y1_or.weight'].T + weights['y1_or.bias'] >= 0).item())
    y1_nand = int((inp @ weights['y1_nand.weight'].T + weights['y1_nand.bias'] >= 0).item())

    # Layer 2
    l2_in = torch.tensor([float(a3), float(a2), float(c2), float(c3), float(y1_or), float(y1_nand)])
    y1 = int((l2_in @ weights['y1.weight'].T + weights['y1.bias'] >= 0).item())
    y2_or = int((l2_in @ weights['y2_or.weight'].T + weights['y2_or.bias'] >= 0).item())
    y2_nand = int((l2_in @ weights['y2_nand.weight'].T + weights['y2_nand.bias'] >= 0).item())
    y3_or = int((l2_in @ weights['y3_or.weight'].T + weights['y3_or.bias'] >= 0).item())
    y3_nand = int((l2_in @ weights['y3_nand.weight'].T + weights['y3_nand.bias'] >= 0).item())

    # Layer 3
    l3_y2 = torch.tensor([float(y2_or), float(y2_nand)])
    l3_y3 = torch.tensor([float(y3_or), float(y3_nand)])
    y2 = int((l3_y2 @ weights['y2.weight'].T + weights['y2.bias'] >= 0).item())
    y3 = int((l3_y3 @ weights['y3.weight'].T + weights['y3.bias'] >= 0).item())

    return [y3, y2, y1, y0]

print("Verifying incrementer4bit...")
errors = 0
for i in range(16):
    a3, a2, a1, a0 = (i >> 3) & 1, (i >> 2) & 1, (i >> 1) & 1, i & 1
    result = incrementer4(a3, a2, a1, a0)
    expected_val = (i + 1) % 16
    expected = [(expected_val >> 3) & 1, (expected_val >> 2) & 1, (expected_val >> 1) & 1, expected_val & 1]
    if result != expected:
        errors += 1
        print(f"ERROR: {i} ({a3}{a2}{a1}{a0}) + 1 = {result}, expected {expected} ({expected_val})")

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}")