Fix evaluator for multi-layer modular and parity circuits

iron_eval.py

@@ -121,10 +121,9 @@ class BatchedFitnessEvaluator:
 
     def _test_twolayer_gate(self, pop: Dict, prefix: str, inputs: torch.Tensor,
                             expected: torch.Tensor) -> torch.Tensor:
-        """Test two-layer gate (XOR, XNOR, BIIMPLIES)."""
+        """Test two-layer gate (XOR, XNOR, BIIMPLIES) - boolean naming (neuron1/neuron2)."""
         pop_size = next(iter(pop.values())).shape[0]
 
-        # Layer 1
         w1_a = pop[f'{prefix}.layer1.neuron1.weight'].view(pop_size, -1)
         b1_a = pop[f'{prefix}.layer1.neuron1.bias'].view(pop_size)
         w1_b = pop[f'{prefix}.layer1.neuron2.weight'].view(pop_size, -1)
@@ -134,7 +133,26 @@ class BatchedFitnessEvaluator:
         h_b = heaviside(inputs @ w1_b.T + b1_b)
         hidden = torch.stack([h_a, h_b], dim=2)
 
-        # Layer 2
+        w2 = pop[f'{prefix}.layer2.weight'].view(pop_size, -1)
+        b2 = pop[f'{prefix}.layer2.bias'].view(pop_size)
+        out = heaviside((hidden * w2.unsqueeze(0)).sum(2) + b2.unsqueeze(0))
+
+        return (out == expected.unsqueeze(1)).float().sum(0)
+
+    def _test_xor_gate_ornand(self, pop: Dict, prefix: str, inputs: torch.Tensor,
+                              expected: torch.Tensor) -> torch.Tensor:
+        """Test two-layer XOR gate - arithmetic naming (or/nand)."""
+        pop_size = next(iter(pop.values())).shape[0]
+
+        w1_or = pop[f'{prefix}.layer1.or.weight'].view(pop_size, -1)
+        b1_or = pop[f'{prefix}.layer1.or.bias'].view(pop_size)
+        w1_nand = pop[f'{prefix}.layer1.nand.weight'].view(pop_size, -1)
+        b1_nand = pop[f'{prefix}.layer1.nand.bias'].view(pop_size)
+
+        h_or = heaviside(inputs @ w1_or.T + b1_or)
+        h_nand = heaviside(inputs @ w1_nand.T + b1_nand)
+        hidden = torch.stack([h_or, h_nand], dim=2)
+
         w2 = pop[f'{prefix}.layer2.weight'].view(pop_size, -1)
         b2 = pop[f'{prefix}.layer2.bias'].view(pop_size)
         out = heaviside((hidden * w2.unsqueeze(0)).sum(2) + b2.unsqueeze(0))
@@ -150,9 +168,8 @@ class BatchedFitnessEvaluator:
         pop_size = next(iter(pop.values())).shape[0]
         scores = torch.zeros(pop_size, device=self.device)
 
-        # Sum (XOR)
-        scores += self._test_twolayer_gate(pop, 'arithmetic.halfadder.sum',
-                                           self.tt2, self.expected['ha_sum'])
+        scores += self._test_xor_gate_ornand(pop, 'arithmetic.halfadder.sum',
+                                             self.tt2, self.expected['ha_sum'])
         # Carry (AND)
         w = pop['arithmetic.halfadder.carry.weight'].view(pop_size, -1)
         b = pop['arithmetic.halfadder.carry.bias'].view(pop_size)
@@ -500,13 +517,46 @@ class BatchedFitnessEvaluator:
     # ERROR DETECTION
     # =========================================================================
 
+    def _eval_xor_gate(self, pop: Dict, prefix: str, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Evaluate XOR gate on batched inputs."""
+        pop_size = next(iter(pop.values())).shape[0]
+
+        w_or = pop[f'{prefix}.layer1.or.weight'].view(pop_size, -1)
+        b_or = pop[f'{prefix}.layer1.or.bias'].view(pop_size)
+        w_nand = pop[f'{prefix}.layer1.nand.weight'].view(pop_size, -1)
+        b_nand = pop[f'{prefix}.layer1.nand.bias'].view(pop_size)
+        w_and = pop[f'{prefix}.layer2.weight'].view(pop_size, -1)
+        b_and = pop[f'{prefix}.layer2.bias'].view(pop_size)
+
+        inp = torch.stack([a, b], dim=2)
+        h_or = heaviside((inp * w_or.unsqueeze(0)).sum(2) + b_or.unsqueeze(0))
+        h_nand = heaviside((inp * w_nand.unsqueeze(0)).sum(2) + b_nand.unsqueeze(0))
+        hidden = torch.stack([h_or, h_nand], dim=2)
+        return heaviside((hidden * w_and.unsqueeze(0)).sum(2) + b_and.unsqueeze(0))
+
     def _test_parity(self, pop: Dict, name: str, even: bool) -> torch.Tensor:
-        """Test parity checker/generator."""
+        """Test parity checker/generator with XOR tree."""
         pop_size = next(iter(pop.values())).shape[0]
-        w = pop[f'error_detection.{name}.weight'].view(pop_size, -1)
-        b = pop[f'error_detection.{name}.bias'].view(pop_size)
+        prefix = f'error_detection.{name}'
+        num_tests = self.test_8bit_bits.shape[0]
+
+        bits = self.test_8bit_bits.unsqueeze(1).expand(-1, pop_size, -1)
+
+        stage1 = []
+        for i, (a, b) in enumerate([(0, 1), (2, 3), (4, 5), (6, 7)]):
+            xor_out = self._eval_xor_gate(pop, f'{prefix}.stage1.xor{i}', bits[:,:,a], bits[:,:,b])
+            stage1.append(xor_out)
+
+        stage2 = []
+        stage2.append(self._eval_xor_gate(pop, f'{prefix}.stage2.xor0', stage1[0], stage1[1]))
+        stage2.append(self._eval_xor_gate(pop, f'{prefix}.stage2.xor1', stage1[2], stage1[3]))
+
+        xor_all = self._eval_xor_gate(pop, f'{prefix}.stage3.xor0', stage2[0], stage2[1])
+
+        w_not = pop[f'{prefix}.output.not.weight'].view(pop_size, -1)
+        b_not = pop[f'{prefix}.output.not.bias'].view(pop_size)
+        out = heaviside(xor_all.unsqueeze(2) * w_not.unsqueeze(0) + b_not.unsqueeze(0)).squeeze(2)
 
-        out = heaviside(self.test_8bit_bits @ w.T + b)
         popcounts = self.test_8bit_bits.sum(1)
         if even:
             expected = ((popcounts.long() % 2) == 0).float()
@@ -519,15 +569,50 @@ class BatchedFitnessEvaluator:
     # MODULAR ARITHMETIC
     # =========================================================================
 
+    def _get_divisible_sums(self, mod: int) -> list:
+        """Get sum values that indicate divisibility by mod."""
+        weights = [(2**(7-i)) % mod for i in range(8)]
+        max_sum = sum(weights)
+        return [k for k in range(0, max_sum + 1) if k % mod == 0]
+
     def _test_modular(self, pop: Dict, mod: int) -> torch.Tensor:
         """Test modular arithmetic circuit."""
         pop_size = next(iter(pop.values())).shape[0]
-        w = pop[f'modular.mod{mod}.weight'].view(pop_size, -1)
-        b = pop[f'modular.mod{mod}.bias'].view(pop_size)
 
-        out = heaviside(self.mod_test_bits @ w.T + b)
-        expected = ((self.mod_test % mod) == 0).float()
+        if mod in [2, 4, 8]:
+            w = pop[f'modular.mod{mod}.weight'].view(pop_size, -1)
+            b = pop[f'modular.mod{mod}.bias'].view(pop_size)
+            out = heaviside(self.mod_test_bits @ w.T + b)
+        else:
+            divisible_sums = self._get_divisible_sums(mod)
+            num_detectors = len(divisible_sums)
+
+            layer1_outputs = []
+            for idx in range(num_detectors):
+                w_geq = pop[f'modular.mod{mod}.layer1.geq{idx}.weight'].view(pop_size, -1)
+                b_geq = pop[f'modular.mod{mod}.layer1.geq{idx}.bias'].view(pop_size)
+                w_leq = pop[f'modular.mod{mod}.layer1.leq{idx}.weight'].view(pop_size, -1)
+                b_leq = pop[f'modular.mod{mod}.layer1.leq{idx}.bias'].view(pop_size)
+
+                geq = heaviside(self.mod_test_bits @ w_geq.T + b_geq)
+                leq = heaviside(self.mod_test_bits @ w_leq.T + b_leq)
+                layer1_outputs.append((geq, leq))
+
+            layer2_outputs = []
+            for idx in range(num_detectors):
+                w_eq = pop[f'modular.mod{mod}.layer2.eq{idx}.weight'].view(pop_size, -1)
+                b_eq = pop[f'modular.mod{mod}.layer2.eq{idx}.bias'].view(pop_size)
+                geq, leq = layer1_outputs[idx]
+                combined = torch.stack([geq, leq], dim=2)
+                eq = heaviside((combined * w_eq.unsqueeze(0)).sum(2) + b_eq.unsqueeze(0))
+                layer2_outputs.append(eq)
+
+            layer2_stack = torch.stack(layer2_outputs, dim=2)
+            w_or = pop[f'modular.mod{mod}.layer3.or.weight'].view(pop_size, -1)
+            b_or = pop[f'modular.mod{mod}.layer3.or.bias'].view(pop_size)
+            out = heaviside((layer2_stack * w_or.unsqueeze(0)).sum(2) + b_or.unsqueeze(0))
 
+        expected = ((self.mod_test % mod) == 0).float()
         return (out == expected.unsqueeze(1)).float().sum(0)
 
     # =========================================================================
@@ -539,35 +624,29 @@ class BatchedFitnessEvaluator:
         pop_size = next(iter(pop.values())).shape[0]
         scores = torch.zeros(pop_size, device=self.device)
 
-        # Test all 8 combinations of (a, b, sel)
         for a in [0, 1]:
             for b in [0, 1]:
                 for sel in [0, 1]:
                     expected = a if sel == 1 else b
 
-                    # MUX uses: and_a, and_b, not_sel, or
                     a_t = torch.full((pop_size,), float(a), device=self.device)
                     b_t = torch.full((pop_size,), float(b), device=self.device)
                     sel_t = torch.full((pop_size,), float(sel), device=self.device)
 
-                    # NOT sel
-                    w_not = pop['combinational.multiplexer2to1.not_sel.weight'].view(pop_size, -1)
-                    b_not = pop['combinational.multiplexer2to1.not_sel.bias'].view(pop_size)
+                    w_not = pop['combinational.multiplexer2to1.not_s.weight'].view(pop_size, -1)
+                    b_not = pop['combinational.multiplexer2to1.not_s.bias'].view(pop_size)
                     not_sel = heaviside(sel_t.unsqueeze(1) @ w_not.T + b_not)
 
-                    # AND(a, sel)
                     inp_a = torch.stack([a_t, sel_t], dim=1)
-                    w_and_a = pop['combinational.multiplexer2to1.and_a.weight'].view(pop_size, -1)
-                    b_and_a = pop['combinational.multiplexer2to1.and_a.bias'].view(pop_size)
+                    w_and_a = pop['combinational.multiplexer2to1.and1.weight'].view(pop_size, -1)
+                    b_and_a = pop['combinational.multiplexer2to1.and1.bias'].view(pop_size)
                     and_a = heaviside((inp_a * w_and_a).sum(1) + b_and_a)
 
-                    # AND(b, not_sel)
                     inp_b = torch.stack([b_t, not_sel.squeeze(1)], dim=1)
-                    w_and_b = pop['combinational.multiplexer2to1.and_b.weight'].view(pop_size, -1)
-                    b_and_b = pop['combinational.multiplexer2to1.and_b.bias'].view(pop_size)
+                    w_and_b = pop['combinational.multiplexer2to1.and0.weight'].view(pop_size, -1)
+                    b_and_b = pop['combinational.multiplexer2to1.and0.bias'].view(pop_size)
                     and_b = heaviside((inp_b * w_and_b).sum(1) + b_and_b)
 
-                    # OR
                     inp_or = torch.stack([and_a, and_b], dim=1)
                     w_or = pop['combinational.multiplexer2to1.or.weight'].view(pop_size, -1)
                     b_or = pop['combinational.multiplexer2to1.or.bias'].view(pop_size)