import torch
from safetensors.torch import save_file

def xor_block(prefix):
    return {
        f'{prefix}.layer1.or.weight': torch.tensor([1.0, 1.0]),
        f'{prefix}.layer1.or.bias': torch.tensor([-1.0]),
        f'{prefix}.layer1.nand.weight': torch.tensor([-1.0, -1.0]),
        f'{prefix}.layer1.nand.bias': torch.tensor([1.0]),
        f'{prefix}.layer2.and.weight': torch.tensor([1.0, 1.0]),
        f'{prefix}.layer2.and.bias': torch.tensor([-2.0]),
    }

weights = {}
weights.update(xor_block('xor1'))
weights.update(xor_block('xor2'))
weights.update(xor_block('xor3'))

save_file(weights, 'model.safetensors')
mag = sum(t.abs().sum().item() for t in weights.values())
print(f'Created model.safetensors (magnitude {mag:.0f})')