import torch
from safetensors.torch import save_file

# Carry Generate: G = a AND b
# Fires when both inputs are 1, guaranteeing a carry

weights = {
    'and.weight': torch.tensor([[1.0, 1.0]], dtype=torch.float32),
    'and.bias': torch.tensor([-2.0], dtype=torch.float32)
}

save_file(weights, 'model.safetensors')

def carry_generate(a, b):
    inp = torch.tensor([float(a), float(b)])
    return int((inp @ weights['and.weight'].T + weights['and.bias'] >= 0).item())

print("Verifying carry-generate...")
errors = 0
for a in [0, 1]:
    for b in [0, 1]:
        result = carry_generate(a, b)
        expected = a & b
        if result != expected:
            errors += 1
            print(f"ERROR: G({a},{b}) = {result}, expected {expected}")

if errors == 0:
    print("All 4 test cases passed!")
else:
    print(f"FAILED: {errors} errors")

print(f"Magnitude: {sum(t.abs().sum().item() for t in weights.values()):.0f}")