phanerozoic's picture
Rename from tiny-RippleCarry2Bit-verified
3579fec verified
"""
Threshold Network for 2-Bit Ripple Carry Adder
Adds two 2-bit numbers with carry in, producing 2-bit sum and carry out.
Built from two cascaded full adders.
"""
import torch
from safetensors.torch import load_file
def heaviside(x):
return (x >= 0).float()
class ThresholdRippleCarry2Bit:
"""
2-bit ripple carry adder: (a1,a0) + (b1,b0) + cin = (cout, s1, s0)
"""
def __init__(self, weights_dict):
self.weights = weights_dict
def half_adder(self, prefix, a, b):
inputs = torch.tensor([float(a), float(b)])
or_out = heaviside((inputs * self.weights[f'{prefix}.sum.layer1.or.weight']).sum() +
self.weights[f'{prefix}.sum.layer1.or.bias'])
nand_out = heaviside((inputs * self.weights[f'{prefix}.sum.layer1.nand.weight']).sum() +
self.weights[f'{prefix}.sum.layer1.nand.bias'])
layer1 = torch.tensor([or_out, nand_out])
sum_out = heaviside((layer1 * self.weights[f'{prefix}.sum.layer2.weight']).sum() +
self.weights[f'{prefix}.sum.layer2.bias'])
carry_out = heaviside((inputs * self.weights[f'{prefix}.carry.weight']).sum() +
self.weights[f'{prefix}.carry.bias'])
return int(sum_out.item()), int(carry_out.item())
def full_adder(self, prefix, a, b, cin):
s1, c1 = self.half_adder(f'{prefix}.ha1', a, b)
sum_out, c2 = self.half_adder(f'{prefix}.ha2', s1, cin)
carry_inputs = torch.tensor([float(c1), float(c2)])
cout = heaviside((carry_inputs * self.weights[f'{prefix}.carry_or.weight']).sum() +
self.weights[f'{prefix}.carry_or.bias'])
return int(sum_out), int(cout.item())
def __call__(self, a0, a1, b0, b1, cin=0):
# FA0: a0 + b0 + cin
s0, c0 = self.full_adder('fa0', a0, b0, cin)
# FA1: a1 + b1 + c0
s1, cout = self.full_adder('fa1', a1, b1, c0)
return s0, s1, cout
@classmethod
def from_safetensors(cls, path="model.safetensors"):
return cls(load_file(path))
if __name__ == "__main__":
model = ThresholdRippleCarry2Bit.from_safetensors("model.safetensors")
print("2-Bit Ripple Carry Adder Tests:")
print("-" * 50)
# Test a few cases
for a in range(4):
for b in range(4):
a0, a1 = a & 1, (a >> 1) & 1
b0, b1 = b & 1, (b >> 1) & 1
s0, s1, cout = model(a0, a1, b0, b1, 0)
result = s0 + 2*s1 + 4*cout
expected = a + b
status = "OK" if result == expected else "FAIL"
print(f"{a} + {b} = {result} (s1s0={s1}{s0}, cout={cout}) [{status}]")