Upload folder using huggingface_hub

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 __pycache__/iron_eval.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text
+__pycache__/iron_eval.cpython-311.pyc filter=lfs diff=lfs merge=lfs -text

__pycache__/iron_eval.cpython-311.pyc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:889b91e75c83db93f3bb79eca5185dcc75309927c4b558944425b30365110603
+size 274934

circuit_llm.py

@@ -0,0 +1,606 @@
+"""
+Circuit-Augmented LLM: Embedding threshold logic circuits into SmolLM2
+======================================================================
+
+Replaces/augments MLP layers with frozen threshold circuits for exact arithmetic.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Tuple
+from safetensors.torch import load_file
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import warnings
+warnings.filterwarnings('ignore')
+
+
+# =============================================================================
+# HEAVISIDE WITH STRAIGHT-THROUGH ESTIMATOR
+# =============================================================================
+
+class HeavisideSTE(torch.autograd.Function):
+    """Heaviside step function with straight-through estimator for backprop."""
+
+    @staticmethod
+    def forward(ctx, x):
+        return (x >= 0).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # STE: pass gradient through unchanged
+        return grad_output
+
+
+def heaviside(x: torch.Tensor) -> torch.Tensor:
+    """Heaviside step: 1 if x >= 0, else 0. Uses STE for training."""
+    return HeavisideSTE.apply(x)
+
+
+# =============================================================================
+# CIRCUIT EXECUTOR - Runs the frozen threshold circuits
+# =============================================================================
+
+class CircuitExecutor(nn.Module):
+    """
+    Executes threshold logic circuits from the safetensors file.
+    All circuit weights are frozen - only interface layers train.
+    """
+
+    def __init__(self, circuit_path: str, device: str = 'cpu'):
+        super().__init__()
+        self.device = device
+
+        # Load all circuit tensors
+        raw_circuits = load_file(circuit_path)
+
+        # Store as frozen parameters (use underscores for valid param names)
+        self.circuits = {}
+        for k, v in raw_circuits.items():
+            safe_name = k.replace('.', '__')
+            self.register_buffer(safe_name, v.float().to(device))
+            self.circuits[k] = safe_name
+
+    def _get(self, name: str) -> torch.Tensor:
+        """Get circuit tensor by original dotted name."""
+        return getattr(self, self.circuits[name])
+
+    # -------------------------------------------------------------------------
+    # Boolean Gates
+    # -------------------------------------------------------------------------
+
+    def eval_and(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """AND gate: output 1 iff both inputs are 1."""
+        inp = torch.stack([a, b], dim=-1)
+        w = self._get('boolean.and.weight')
+        bias = self._get('boolean.and.bias')
+        return heaviside(inp @ w + bias)
+
+    def eval_or(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """OR gate: output 1 if either input is 1."""
+        inp = torch.stack([a, b], dim=-1)
+        w = self._get('boolean.or.weight')
+        bias = self._get('boolean.or.bias')
+        return heaviside(inp @ w + bias)
+
+    def eval_xor(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """XOR gate: two-layer network (not linearly separable)."""
+        inp = torch.stack([a, b], dim=-1)
+
+        # Layer 1: OR and NAND neurons
+        w1_n1 = self._get('boolean.xor.layer1.neuron1.weight')
+        b1_n1 = self._get('boolean.xor.layer1.neuron1.bias')
+        w1_n2 = self._get('boolean.xor.layer1.neuron2.weight')
+        b1_n2 = self._get('boolean.xor.layer1.neuron2.bias')
+
+        h1 = heaviside(inp @ w1_n1 + b1_n1)
+        h2 = heaviside(inp @ w1_n2 + b1_n2)
+        hidden = torch.stack([h1, h2], dim=-1)
+
+        # Layer 2: AND of hidden
+        w2 = self._get('boolean.xor.layer2.weight')
+        b2 = self._get('boolean.xor.layer2.bias')
+
+        return heaviside(hidden @ w2 + b2)
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: Full Adder
+    # -------------------------------------------------------------------------
+
+    def eval_full_adder(self, a: torch.Tensor, b: torch.Tensor,
+                        cin: torch.Tensor, prefix: str) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Full adder: sum = a XOR b XOR cin, cout = (a AND b) OR (cin AND (a XOR b))
+        Returns (sum_bit, carry_out)
+        """
+        inp_ab = torch.stack([a, b], dim=-1)
+
+        # HA1: a XOR b
+        w1_or = self._get(f'{prefix}.ha1.sum.layer1.or.weight')
+        b1_or = self._get(f'{prefix}.ha1.sum.layer1.or.bias')
+        w1_nand = self._get(f'{prefix}.ha1.sum.layer1.nand.weight')
+        b1_nand = self._get(f'{prefix}.ha1.sum.layer1.nand.bias')
+        w2 = self._get(f'{prefix}.ha1.sum.layer2.weight')
+        b2 = self._get(f'{prefix}.ha1.sum.layer2.bias')
+
+        h_or = heaviside(inp_ab @ w1_or + b1_or)
+        h_nand = heaviside(inp_ab @ w1_nand + b1_nand)
+        hidden = torch.stack([h_or, h_nand], dim=-1)
+        ha1_sum = heaviside(hidden @ w2 + b2)
+
+        # HA1 carry
+        w_c1 = self._get(f'{prefix}.ha1.carry.weight')
+        b_c1 = self._get(f'{prefix}.ha1.carry.bias')
+        ha1_carry = heaviside(inp_ab @ w_c1 + b_c1)
+
+        # HA2: ha1_sum XOR cin
+        inp_ha2 = torch.stack([ha1_sum, cin], dim=-1)
+        w1_or = self._get(f'{prefix}.ha2.sum.layer1.or.weight')
+        b1_or = self._get(f'{prefix}.ha2.sum.layer1.or.bias')
+        w1_nand = self._get(f'{prefix}.ha2.sum.layer1.nand.weight')
+        b1_nand = self._get(f'{prefix}.ha2.sum.layer1.nand.bias')
+        w2 = self._get(f'{prefix}.ha2.sum.layer2.weight')
+        b2 = self._get(f'{prefix}.ha2.sum.layer2.bias')
+
+        h_or = heaviside(inp_ha2 @ w1_or + b1_or)
+        h_nand = heaviside(inp_ha2 @ w1_nand + b1_nand)
+        hidden = torch.stack([h_or, h_nand], dim=-1)
+        ha2_sum = heaviside(hidden @ w2 + b2)
+
+        # HA2 carry
+        w_c2 = self._get(f'{prefix}.ha2.carry.weight')
+        b_c2 = self._get(f'{prefix}.ha2.carry.bias')
+        ha2_carry = heaviside(inp_ha2 @ w_c2 + b_c2)
+
+        # Carry out = ha1_carry OR ha2_carry
+        inp_cout = torch.stack([ha1_carry, ha2_carry], dim=-1)
+        w_or = self._get(f'{prefix}.carry_or.weight')
+        b_or = self._get(f'{prefix}.carry_or.bias')
+        cout = heaviside(inp_cout @ w_or + b_or)
+
+        return ha2_sum, cout
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: 8-bit Ripple Carry Adder
+    # -------------------------------------------------------------------------
+
+    def add_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        8-bit ripple carry addition.
+        a_bits, b_bits: [..., 8] tensors of bits (LSB first)
+        Returns: (result_bits [..., 8], carry_out [...])
+        """
+        batch_shape = a_bits.shape[:-1]
+        carry = torch.zeros(batch_shape, device=a_bits.device)
+        result_bits = []
+
+        for i in range(8):
+            a_i = a_bits[..., i]
+            b_i = b_bits[..., i]
+            sum_bit, carry = self.eval_full_adder(
+                a_i, b_i, carry,
+                f'arithmetic.ripplecarry8bit.fa{i}'
+            )
+            result_bits.append(sum_bit)
+
+        return torch.stack(result_bits, dim=-1), carry
+
+    # -------------------------------------------------------------------------
+    # Arithmetic: 8-bit Comparators
+    # -------------------------------------------------------------------------
+
+    def greater_than_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a > b, else 0. Bits are MSB first."""
+        diff = a_bits - b_bits  # [..., 8]
+        w = self._get('arithmetic.greaterthan8bit.comparator')
+        score = (diff * w).sum(dim=-1)
+        return (score > 0).float()
+
+    def less_than_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a < b, else 0. Bits are MSB first."""
+        diff = b_bits - a_bits  # [..., 8]
+        w = self._get('arithmetic.lessthan8bit.comparator')
+        score = (diff * w).sum(dim=-1)
+        return (score > 0).float()
+
+    def equal_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
+        """Returns 1 if a == b, else 0."""
+        gt = self.greater_than_8bit(a_bits, b_bits)
+        lt = self.less_than_8bit(a_bits, b_bits)
+        return (1 - gt) * (1 - lt)
+
+
+# =============================================================================
+# BIT EXTRACTION / INJECTION INTERFACES
+# =============================================================================
+
+class BitExtractor(nn.Module):
+    """
+    Learns to extract 8-bit operands from token embeddings.
+    Maps embedding -> 16 bits (two 8-bit operands).
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+        # Project to logits, then binarize
+        self.proj = nn.Linear(d_model, 16)
+
+        # Learnable temperature for sigmoid approximation during training
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        x: [..., d_model]
+        Returns: a_bits [..., 8], b_bits [..., 8] (LSB first for arithmetic)
+        """
+        logits = self.proj(x)  # [..., 16]
+
+        # Binarize with STE
+        bits = heaviside(logits)
+
+        # Split into two operands
+        a_bits = bits[..., :8]
+        b_bits = bits[..., 8:]
+
+        return a_bits, b_bits
+
+
+class BitInjector(nn.Module):
+    """
+    Learns to inject circuit results back into embedding space.
+    Maps 16 bits (result + flags) -> embedding delta.
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.d_model = d_model
+
+        # Project bits to embedding
+        self.proj = nn.Linear(16, d_model)
+
+        # Learnable scale
+        self.scale = nn.Parameter(torch.tensor(0.1))
+
+    def forward(self, result_bits: torch.Tensor, flags: torch.Tensor) -> torch.Tensor:
+        """
+        result_bits: [..., 8]
+        flags: [..., 8] (carry, overflow, zero, negative, etc.)
+        Returns: [..., d_model]
+        """
+        combined = torch.cat([result_bits, flags], dim=-1)  # [..., 16]
+        return self.proj(combined) * self.scale
+
+
+# =============================================================================
+# CIRCUIT-AUGMENTED MLP BLOCK
+# =============================================================================
+
+class CircuitAugmentedMLP(nn.Module):
+    """
+    MLP block augmented with frozen threshold circuits.
+
+    The original MLP path runs in parallel with the circuit path.
+    A learned router decides how much to use each.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        intermediate_size: int,
+        circuit_path: str,
+        device: str = 'cpu'
+    ):
+        super().__init__()
+        self.d_model = d_model
+
+        # Original MLP components (will be loaded from pretrained)
+        self.gate_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, d_model, bias=False)
+        self.act_fn = nn.SiLU()
+
+        # Circuit components
+        self.circuits = CircuitExecutor(circuit_path, device)
+        self.bit_extractor = BitExtractor(d_model)
+        self.bit_injector = BitInjector(d_model)
+
+        # Router: decides circuit vs MLP contribution
+        self.router = nn.Sequential(
+            nn.Linear(d_model, 64),
+            nn.ReLU(),
+            nn.Linear(64, 2),
+            nn.Softmax(dim=-1)
+        )
+
+        # Operation selector (which arithmetic op to perform)
+        self.op_selector = nn.Sequential(
+            nn.Linear(d_model, 32),
+            nn.ReLU(),
+            nn.Linear(32, 4),  # add, sub, compare, passthrough
+            nn.Softmax(dim=-1)
+        )
+
+    def _compute_flags(self, result_bits: torch.Tensor, carry: torch.Tensor) -> torch.Tensor:
+        """Compute status flags from result."""
+        batch_shape = result_bits.shape[:-1]
+
+        # Zero flag: all bits are 0
+        zero = (result_bits.sum(dim=-1) == 0).float()
+
+        # Negative flag: MSB is 1 (two's complement)
+        negative = result_bits[..., 7]
+
+        # Carry flag
+        carry_flag = carry
+
+        # Pad to 8 flags
+        flags = torch.zeros(*batch_shape, 8, device=result_bits.device)
+        flags[..., 0] = zero
+        flags[..., 1] = negative
+        flags[..., 2] = carry_flag
+
+        return flags
+
+    def _circuit_forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run input through threshold circuits."""
+        # Extract operands
+        a_bits, b_bits = self.bit_extractor(x)
+
+        # Get operation weights
+        op_weights = self.op_selector(x)  # [..., 4]
+
+        # Compute addition
+        add_result, add_carry = self.circuits.add_8bit(a_bits, b_bits)
+        add_flags = self._compute_flags(add_result, add_carry)
+
+        # Compute subtraction (a + (~b) + 1, simplified: just use add for now)
+        # For MVP, we'll focus on addition
+
+        # Inject result back
+        circuit_delta = self.bit_injector(add_result, add_flags)
+
+        return circuit_delta
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: [batch, seq_len, d_model]
+        Returns: [batch, seq_len, d_model]
+        """
+        # Original MLP path
+        mlp_out = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+        # Circuit path
+        circuit_out = self._circuit_forward(x)
+
+        # Route between paths
+        route_weights = self.router(x)  # [..., 2]
+        mlp_weight = route_weights[..., 0:1]
+        circuit_weight = route_weights[..., 1:2]
+
+        # Combine: MLP output + weighted circuit contribution
+        output = mlp_out + circuit_weight * circuit_out
+
+        return output
+
+
+# =============================================================================
+# MODEL SURGERY: Insert circuits into SmolLM2
+# =============================================================================
+
+def augment_smollm2_with_circuits(
+    model: AutoModelForCausalLM,
+    circuit_path: str,
+    layer_indices: list = None,
+    device: str = 'cpu'
+) -> AutoModelForCausalLM:
+    """
+    Surgically insert circuit blocks into SmolLM2's MLP layers.
+
+    Args:
+        model: Pretrained SmolLM2 model
+        circuit_path: Path to neural_computer.safetensors
+        layer_indices: Which layers to augment (default: middle layers)
+        device: Device for circuit tensors
+
+    Returns:
+        Modified model with circuit-augmented MLPs
+    """
+    config = model.config
+    num_layers = config.num_hidden_layers
+
+    # Default: augment middle third of layers
+    if layer_indices is None:
+        start = num_layers // 3
+        end = 2 * num_layers // 3
+        layer_indices = list(range(start, end))
+
+    print(f"Augmenting layers {layer_indices} with threshold circuits...")
+
+    for idx in layer_indices:
+        layer = model.model.layers[idx]
+        old_mlp = layer.mlp
+
+        # Create augmented MLP
+        new_mlp = CircuitAugmentedMLP(
+            d_model=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            circuit_path=circuit_path,
+            device=device
+        )
+
+        # Copy pretrained weights
+        new_mlp.gate_proj.weight.data = old_mlp.gate_proj.weight.data.clone()
+        new_mlp.up_proj.weight.data = old_mlp.up_proj.weight.data.clone()
+        new_mlp.down_proj.weight.data = old_mlp.down_proj.weight.data.clone()
+
+        # Replace
+        layer.mlp = new_mlp
+
+    # Freeze circuit weights, keep interfaces trainable
+    for name, param in model.named_parameters():
+        if 'circuits' in name:
+            param.requires_grad = False
+
+    print(f"Done. Circuit weights frozen, interfaces trainable.")
+
+    return model
+
+
+# =============================================================================
+# TRAINING UTILITIES
+# =============================================================================
+
+def generate_arithmetic_batch(batch_size: int, max_val: int = 255) -> Tuple[list, list]:
+    """Generate batch of arithmetic problems and solutions."""
+    prompts = []
+    targets = []
+
+    for _ in range(batch_size):
+        a = torch.randint(0, max_val + 1, (1,)).item()
+        b = torch.randint(0, max_val + 1, (1,)).item()
+        result = (a + b) % 256
+
+        prompts.append(f"{a} + {b} =")
+        targets.append(f" {result}")
+
+    return prompts, targets
+
+
+def evaluate_arithmetic(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    n_problems: int = 100,
+    device: str = 'cpu'
+) -> dict:
+    """Evaluate model on random arithmetic problems."""
+    correct = 0
+    total = 0
+    errors = []
+
+    model.eval()
+
+    for _ in range(n_problems):
+        a = torch.randint(0, 256, (1,)).item()
+        b = torch.randint(0, 256, (1,)).item()
+        expected = (a + b) % 256
+
+        prompt = f"{a} + {b} ="
+        inputs = tokenizer(prompt, return_tensors='pt').to(device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=10,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Extract number from response
+        try:
+            # Find the part after "="
+            answer_part = response.split('=')[-1].strip()
+            # Extract first number
+            predicted = int(''.join(c for c in answer_part.split()[0] if c.isdigit()))
+
+            if predicted == expected:
+                correct += 1
+            else:
+                errors.append((a, b, expected, predicted))
+        except:
+            errors.append((a, b, expected, "parse_error"))
+
+        total += 1
+
+    return {
+        'accuracy': correct / total,
+        'correct': correct,
+        'total': total,
+        'errors': errors[:10]  # First 10 errors
+    }
+
+
+# =============================================================================
+# MAIN: Demo
+# =============================================================================
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Circuit-Augmented LLM Demo')
+    parser.add_argument('--circuit-path', type=str,
+                        default='./neural_computer.safetensors',
+                        help='Path to circuit weights')
+    parser.add_argument('--device', type=str, default='cpu',
+                        help='Device (cpu or cuda)')
+    parser.add_argument('--eval-only', action='store_true',
+                        help='Only evaluate, do not augment')
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print(" CIRCUIT-AUGMENTED LLM")
+    print("=" * 70)
+
+    # Load tokenizer and model
+    print("\n[1] Loading SmolLM2-360M...")
+    model_id = "HuggingFaceTB/SmolLM2-360M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+
+    print(f"    Parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Baseline evaluation
+    print("\n[2] Baseline arithmetic evaluation...")
+    baseline = evaluate_arithmetic(model, tokenizer, n_problems=50, device=args.device)
+    print(f"    Accuracy: {baseline['accuracy']*100:.1f}% ({baseline['correct']}/{baseline['total']})")
+    if baseline['errors']:
+        print(f"    Sample errors:")
+        for a, b, exp, got in baseline['errors'][:5]:
+            print(f"      {a} + {b} = {exp}, model said {got}")
+
+    if args.eval_only:
+        print("\nDone (eval only mode).")
+        exit(0)
+
+    # Augment with circuits
+    print(f"\n[3] Augmenting with threshold circuits...")
+    print(f"    Circuit path: {args.circuit_path}")
+    model = augment_smollm2_with_circuits(
+        model,
+        args.circuit_path,
+        device=args.device
+    )
+
+    new_params = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"    Total parameters: {new_params:,}")
+    print(f"    Trainable parameters: {trainable:,}")
+
+    # Test circuit execution directly
+    print("\n[4] Testing circuit execution...")
+    circuit_exec = CircuitExecutor(args.circuit_path, args.device)
+
+    test_cases = [(127, 128), (255, 1), (0, 0), (100, 55)]
+    for a, b in test_cases:
+        # Convert to bits (LSB first)
+        a_bits = torch.tensor([(a >> i) & 1 for i in range(8)], dtype=torch.float32)
+        b_bits = torch.tensor([(b >> i) & 1 for i in range(8)], dtype=torch.float32)
+
+        result_bits, carry = circuit_exec.add_8bit(
+            a_bits.unsqueeze(0),
+            b_bits.unsqueeze(0)
+        )
+
+        # Convert result bits back to int
+        result = sum(int(result_bits[0, i].item()) * (2**i) for i in range(8))
+        expected = (a + b) % 256
+
+        status = "OK" if result == expected else "FAIL"
+        print(f"    {a} + {b} = {result} (expected {expected}) [{status}]")
+
+    print("\n[5] Model ready for fine-tuning.")
+    print("    Next: Train interface layers on arithmetic examples.")
+    print("=" * 70)

stress_test.py

@@ -0,0 +1,367 @@
+"""
+WILD STRESS TESTS - Push the threshold CPU to its limits
+"""
+import torch
+from safetensors.torch import load_file
+
+model = load_file('./neural_computer.safetensors')
+model = {k: v.float() for k, v in model.items()}
+
+def heaviside(x):
+    return (x >= 0).float()
+
+def int_to_bits(val, width=8):
+    return torch.tensor([(val >> (width-1-i)) & 1 for i in range(width)], dtype=torch.float32)
+
+def bits_to_int(bits):
+    val = 0
+    for i, b in enumerate(bits):
+        val |= (int(b.item()) << (len(bits)-1-i))
+    return val
+
+# === BASIC PRIMITIVES ===
+
+def eval_xor(a, b):
+    inp = torch.tensor([float(a), float(b)], dtype=torch.float32)
+    w1_n1 = model['boolean.xor.layer1.neuron1.weight']
+    b1_n1 = model['boolean.xor.layer1.neuron1.bias']
+    w1_n2 = model['boolean.xor.layer1.neuron2.weight']
+    b1_n2 = model['boolean.xor.layer1.neuron2.bias']
+    w2 = model['boolean.xor.layer2.weight']
+    b2 = model['boolean.xor.layer2.bias']
+    h1 = heaviside(inp @ w1_n1 + b1_n1)
+    h2 = heaviside(inp @ w1_n2 + b1_n2)
+    hidden = torch.tensor([h1.item(), h2.item()])
+    return int(heaviside(hidden @ w2 + b2).item())
+
+def eval_and(a, b):
+    inp = torch.tensor([float(a), float(b)], dtype=torch.float32)
+    return int(heaviside(inp @ model['boolean.and.weight'] + model['boolean.and.bias']).item())
+
+def eval_or(a, b):
+    inp = torch.tensor([float(a), float(b)], dtype=torch.float32)
+    return int(heaviside(inp @ model['boolean.or.weight'] + model['boolean.or.bias']).item())
+
+def eval_not(a):
+    inp = torch.tensor([float(a)], dtype=torch.float32)
+    return int(heaviside(inp @ model['boolean.not.weight'] + model['boolean.not.bias']).item())
+
+def eval_xor_arith(inp, prefix):
+    w1_or = model[f'{prefix}.layer1.or.weight']
+    b1_or = model[f'{prefix}.layer1.or.bias']
+    w1_nand = model[f'{prefix}.layer1.nand.weight']
+    b1_nand = model[f'{prefix}.layer1.nand.bias']
+    w2 = model[f'{prefix}.layer2.weight']
+    b2 = model[f'{prefix}.layer2.bias']
+    h_or = heaviside(inp @ w1_or + b1_or)
+    h_nand = heaviside(inp @ w1_nand + b1_nand)
+    hidden = torch.tensor([h_or.item(), h_nand.item()])
+    return heaviside(hidden @ w2 + b2).item()
+
+def eval_full_adder(a, b, cin, prefix):
+    inp_ab = torch.tensor([a, b], dtype=torch.float32)
+    ha1_sum = eval_xor_arith(inp_ab, f'{prefix}.ha1.sum')
+    ha1_carry = heaviside(inp_ab @ model[f'{prefix}.ha1.carry.weight'] + model[f'{prefix}.ha1.carry.bias']).item()
+    inp_ha2 = torch.tensor([ha1_sum, cin], dtype=torch.float32)
+    ha2_sum = eval_xor_arith(inp_ha2, f'{prefix}.ha2.sum')
+    ha2_carry = heaviside(inp_ha2 @ model[f'{prefix}.ha2.carry.weight'] + model[f'{prefix}.ha2.carry.bias']).item()
+    inp_cout = torch.tensor([ha1_carry, ha2_carry], dtype=torch.float32)
+    cout = heaviside(inp_cout @ model[f'{prefix}.carry_or.weight'] + model[f'{prefix}.carry_or.bias']).item()
+    return int(ha2_sum), int(cout)
+
+def add_8bit(a, b):
+    carry = 0.0
+    result = 0
+    for i in range(8):
+        s, carry = eval_full_adder(float((a >> i) & 1), float((b >> i) & 1), carry, f'arithmetic.ripplecarry8bit.fa{i}')
+        result |= (s << i)
+    return result, int(carry)
+
+def sub_8bit(a, b):
+    # a - b = a + (~b + 1)
+    not_b = 0
+    for i in range(8):
+        not_b |= (eval_not((b >> i) & 1) << i)
+    temp, _ = add_8bit(a, not_b)
+    result, _ = add_8bit(temp, 1)
+    return result
+
+def gt(a, b):
+    a_bits, b_bits = int_to_bits(a), int_to_bits(b)
+    w = model['arithmetic.greaterthan8bit.comparator']
+    return 1 if ((a_bits - b_bits) @ w).item() > 0 else 0
+
+def lt(a, b):
+    a_bits, b_bits = int_to_bits(a), int_to_bits(b)
+    w = model['arithmetic.lessthan8bit.comparator']
+    return 1 if ((b_bits - a_bits) @ w).item() > 0 else 0
+
+def eq(a, b):
+    return 1 if (gt(a,b) == 0 and lt(a,b) == 0) else 0
+
+def popcount(val):
+    bits = int_to_bits(val)
+    w = model['pattern_recognition.popcount.weight']
+    b = model['pattern_recognition.popcount.bias']
+    return int((bits @ w + b).item())
+
+print('='*70)
+print('WILD STRESS TESTS')
+print('='*70)
+
+# === TEST 1: FACTORIAL ===
+print('\n[1] FACTORIAL via chained multiply-add')
+def factorial(n):
+    result = 1
+    for i in range(2, n+1):
+        new_result = 0
+        for _ in range(i):
+            new_result, _ = add_8bit(new_result, result)
+            new_result &= 0xFF
+        result = new_result
+    return result
+
+for n in [1, 2, 3, 4, 5]:
+    got = factorial(n)
+    expected = [1, 1, 2, 6, 24, 120][n]
+    status = 'OK' if got == expected else 'FAIL'
+    print(f'  {n}! = {got} (expected {expected}) [{status}]')
+
+# === TEST 2: GCD ===
+print('\n[2] GCD via Euclidean algorithm')
+def gcd(a, b):
+    iterations = 0
+    while not eq(b, 0) and iterations < 100:
+        temp = a
+        while not lt(temp, b) and not eq(temp, 0) and iterations < 100:
+            temp = sub_8bit(temp, b)
+            iterations += 1
+        a, b = b, temp
+        iterations += 1
+    return a
+
+test_gcds = [(48, 18, 6), (100, 35, 5), (252, 105, 21), (17, 13, 1), (128, 64, 64)]
+for a, b, expected in test_gcds:
+    got = gcd(a, b)
+    status = 'OK' if got == expected else 'FAIL'
+    print(f'  gcd({a}, {b}) = {got} (expected {expected}) [{status}]')
+
+# === TEST 3: FIBONACCI ===
+print('\n[3] FIBONACCI until overflow')
+def fib_sequence():
+    a, b = 0, 1
+    seq = [a, b]
+    for _ in range(20):
+        next_val, carry = add_8bit(a, b)
+        if carry:
+            break
+        seq.append(next_val)
+        a, b = b, next_val
+    return seq
+
+fib = fib_sequence()
+expected_fib = [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233]
+print(f'  Computed: {fib[:len(expected_fib)]}')
+print(f'  Expected: {expected_fib}')
+print(f'  Match: {fib[:len(expected_fib)] == expected_fib}')
+
+# === TEST 4: PRIME CHECK ===
+print('\n[4] PRIME CHECK via trial division')
+def is_prime(n):
+    if n < 2: return False
+    if n == 2: return True
+    if (n & 1) == 0: return False
+
+    i = 3
+    while i * i <= n and i < n:
+        temp = n
+        while temp >= i:
+            temp = sub_8bit(temp, i)
+        if eq(temp, 0):
+            return False
+        i += 2
+    return True
+
+primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
+non_primes = [4, 6, 8, 9, 10, 12, 14, 15, 16, 18]
+
+prime_pass = sum(1 for p in primes if is_prime(p))
+non_prime_pass = sum(1 for n in non_primes if not is_prime(n))
+print(f'  Primes correctly identified: {prime_pass}/10')
+print(f'  Non-primes correctly rejected: {non_prime_pass}/10')
+
+# === TEST 5: INTEGER SQRT ===
+print('\n[5] INTEGER SQUARE ROOT via binary search')
+def isqrt(n):
+    if n == 0: return 0
+    lo, hi = 1, min(n, 15)  # limit for 8-bit
+    result = 0
+    iterations = 0
+    while lo <= hi and iterations < 50:
+        mid = (lo + hi) >> 1
+        sq = 0
+        for _ in range(mid):
+            sq, _ = add_8bit(sq, mid)
+            sq &= 0xFF
+
+        if sq <= n:
+            result = mid
+            lo = mid + 1
+        else:
+            hi = mid - 1
+        iterations += 1
+    return result
+
+sqrt_tests = [(0, 0), (1, 1), (4, 2), (9, 3), (16, 4), (25, 5), (36, 6), (49, 7), (64, 8), (81, 9), (100, 10), (144, 12)]
+sqrt_pass = 0
+for n, expected in sqrt_tests:
+    got = isqrt(n)
+    if got == expected:
+        sqrt_pass += 1
+print(f'  Passed: {sqrt_pass}/{len(sqrt_tests)}')
+
+# === TEST 6: COLLATZ ===
+print('\n[6] COLLATZ CONJECTURE iterations')
+def collatz_steps(n):
+    steps = 0
+    while n != 1 and steps < 200:
+        if (n & 1) == 0:
+            n = n >> 1
+        else:
+            temp, _ = add_8bit(n, n)
+            temp, _ = add_8bit(temp, n)
+            n, _ = add_8bit(temp, 1)
+            n &= 0xFF
+        steps += 1
+        if n == 0: break
+    return steps
+
+collatz_tests = [(1, 0), (2, 1), (3, 7), (6, 8)]
+for start, expected in collatz_tests:
+    got = collatz_steps(start)
+    status = 'OK' if got == expected else f'got {got}'
+    print(f'  collatz({start}) = {got} steps [{status}]')
+
+# === TEST 7: SORT BY POPCOUNT ===
+print('\n[7] SORT BY HAMMING WEIGHT (popcount)')
+values = [0b11111111, 0b00000001, 0b10101010, 0b00001111, 0b11110000, 0b00000000]
+weighted = [(v, popcount(v)) for v in values]
+for i in range(len(weighted)):
+    for j in range(len(weighted) - 1):
+        if gt(weighted[j][1], weighted[j+1][1]):
+            weighted[j], weighted[j+1] = weighted[j+1], weighted[j]
+
+print(f'  Sorted by popcount:')
+for v, p in weighted:
+    print(f'    {bin(v):>12} -> popcount = {p}')
+
+# === TEST 8: XOR CHECKSUM ===
+print('\n[8] XOR CHECKSUM of message')
+message = [0x48, 0x65, 0x6C, 0x6C, 0x6F]  # "Hello"
+checksum = 0
+for byte in message:
+    for i in range(8):
+        bit_a = (checksum >> i) & 1
+        bit_b = (byte >> i) & 1
+        xor_bit = eval_xor(bit_a, bit_b)
+        checksum = (checksum & ~(1 << i)) | (xor_bit << i)
+
+expected_checksum = 0x48 ^ 0x65 ^ 0x6C ^ 0x6C ^ 0x6F
+status = 'OK' if checksum == expected_checksum else 'FAIL'
+print(f'  Message: {[hex(b) for b in message]}')
+print(f'  XOR checksum: {hex(checksum)} (expected {hex(expected_checksum)}) [{status}]')
+
+# === TEST 9: PARITY TREE ===
+print('\n[9] 8-BIT PARITY (full XOR tree)')
+def parity_8bit(val):
+    bits = [(val >> i) & 1 for i in range(8)]
+    s1 = [eval_xor(bits[0], bits[1]), eval_xor(bits[2], bits[3]),
+          eval_xor(bits[4], bits[5]), eval_xor(bits[6], bits[7])]
+    s2 = [eval_xor(s1[0], s1[1]), eval_xor(s1[2], s1[3])]
+    return eval_xor(s2[0], s2[1])
+
+parity_tests = [(0x00, 0), (0xFF, 0), (0x01, 1), (0x03, 0), (0x07, 1), (0xAA, 0), (0x55, 0), (0x81, 0), (0x80, 1)]
+parity_pass = sum(1 for v, exp in parity_tests if parity_8bit(v) == exp)
+print(f'  Passed: {parity_pass}/{len(parity_tests)}')
+
+# === TEST 10: OVERFLOW CASCADE ===
+print('\n[10] OVERFLOW CASCADE (255 + 1 chain)')
+val = 255
+carries = []
+for i in range(5):
+    val, carry = add_8bit(val, 1)
+    carries.append(carry)
+print(f'  255 -> +1 -> +1 -> +1 -> +1 -> +1')
+print(f'  Carries: {carries}')
+print(f'  Final value: {val} (expected 4) [{"OK" if val == 4 else "FAIL"}]')
+
+# === TEST 11: POWER OF 2 CHECK ===
+print('\n[11] POWER OF 2 detection (popcount == 1)')
+def is_power_of_2(n):
+    if n == 0: return False
+    return popcount(n) == 1
+
+pow2_tests = [(1, True), (2, True), (4, True), (8, True), (16, True), (32, True), (64, True), (128, True),
+              (3, False), (5, False), (6, False), (7, False), (9, False), (15, False), (255, False)]
+pow2_pass = sum(1 for n, exp in pow2_tests if is_power_of_2(n) == exp)
+print(f'  Passed: {pow2_pass}/{len(pow2_tests)}')
+
+# === TEST 12: BYTE REVERSE ===
+print('\n[12] BYTE REVERSE via bit manipulation')
+def reverse_bits(val):
+    result = 0
+    for i in range(8):
+        bit = (val >> i) & 1
+        result |= (bit << (7 - i))
+    return result
+
+reverse_tests = [(0b10000000, 0b00000001), (0b11110000, 0b00001111), (0b10101010, 0b01010101), (0b00000000, 0b00000000), (0b11111111, 0b11111111)]
+reverse_pass = sum(1 for inp, exp in reverse_tests if reverse_bits(inp) == exp)
+print(f'  Passed: {reverse_pass}/{len(reverse_tests)}')
+
+# === TEST 13: MAX/MIN via comparator ===
+print('\n[13] MAX and MIN of array')
+def find_max(arr):
+    m = arr[0]
+    for x in arr[1:]:
+        if gt(x, m):
+            m = x
+    return m
+
+def find_min(arr):
+    m = arr[0]
+    for x in arr[1:]:
+        if lt(x, m):
+            m = x
+    return m
+
+test_arr = [42, 17, 255, 0, 128, 64, 33]
+got_max = find_max(test_arr)
+got_min = find_min(test_arr)
+print(f'  Array: {test_arr}')
+print(f'  Max: {got_max} (expected 255) [{"OK" if got_max == 255 else "FAIL"}]')
+print(f'  Min: {got_min} (expected 0) [{"OK" if got_min == 0 else "FAIL"}]')
+
+# === TEST 14: LFSR (pseudo-random) ===
+print('\n[14] 8-BIT LFSR (taps at 8,6,5,4)')
+def lfsr_step(state):
+    # Taps: 8, 6, 5, 4 (for maximal length)
+    bit = eval_xor((state >> 0) & 1, (state >> 2) & 1)
+    bit = eval_xor(bit, (state >> 3) & 1)
+    bit = eval_xor(bit, (state >> 4) & 1)
+    return ((state >> 1) | (bit << 7)) & 0xFF
+
+state = 1
+seen = set()
+for i in range(300):
+    if state in seen:
+        break
+    seen.add(state)
+    state = lfsr_step(state)
+
+print(f'  Period: {len(seen)} (max possible: 255)')
+print(f'  Full period: {"OK" if len(seen) == 255 else "FAIL"}')
+
+print('\n' + '='*70)
+print('STRESS TESTS COMPLETE')
+print('='*70)

train_circuit_interface.py

@@ -0,0 +1,306 @@
+"""
+Train the circuit interface layers on arithmetic examples.
+============================================================
+
+The threshold circuits are frozen - we only train:
+- BitExtractor: embedding -> operand bits
+- BitInjector: result bits -> embedding
+- Router: when to use circuits vs MLP
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from tqdm import tqdm
+import argparse
+import warnings
+warnings.filterwarnings('ignore')
+
+from circuit_llm import (
+    augment_smollm2_with_circuits,
+    evaluate_arithmetic,
+    CircuitExecutor
+)
+
+
+# =============================================================================
+# ARITHMETIC DATASET
+# =============================================================================
+
+class ArithmeticDataset(Dataset):
+    """Dataset of 8-bit addition problems."""
+
+    def __init__(self, tokenizer, n_samples: int = 10000, max_val: int = 255):
+        self.tokenizer = tokenizer
+        self.n_samples = n_samples
+        self.max_val = max_val
+
+        # Pre-generate all examples
+        self.examples = []
+        for _ in range(n_samples):
+            a = torch.randint(0, max_val + 1, (1,)).item()
+            b = torch.randint(0, max_val + 1, (1,)).item()
+            result = (a + b) % 256
+
+            prompt = f"{a} + {b} ="
+            target = f" {result}"
+
+            self.examples.append((prompt, target, a, b, result))
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        prompt, target, a, b, result = self.examples[idx]
+
+        # Tokenize
+        prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+        target_ids = self.tokenizer.encode(target, add_special_tokens=False)
+
+        input_ids = prompt_ids + target_ids
+        labels = [-100] * len(prompt_ids) + target_ids  # Only predict target
+
+        return {
+            'input_ids': torch.tensor(input_ids),
+            'labels': torch.tensor(labels),
+            'a': a,
+            'b': b,
+            'result': result
+        }
+
+
+def collate_fn(batch):
+    """Collate with padding."""
+    max_len = max(len(item['input_ids']) for item in batch)
+
+    input_ids = []
+    labels = []
+    attention_mask = []
+
+    for item in batch:
+        pad_len = max_len - len(item['input_ids'])
+
+        input_ids.append(
+            torch.cat([item['input_ids'], torch.zeros(pad_len, dtype=torch.long)])
+        )
+        labels.append(
+            torch.cat([item['labels'], torch.full((pad_len,), -100, dtype=torch.long)])
+        )
+        attention_mask.append(
+            torch.cat([torch.ones(len(item['input_ids'])), torch.zeros(pad_len)])
+        )
+
+    return {
+        'input_ids': torch.stack(input_ids),
+        'labels': torch.stack(labels),
+        'attention_mask': torch.stack(attention_mask),
+    }
+
+
+# =============================================================================
+# TRAINING LOOP
+# =============================================================================
+
+def train_interface(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    n_epochs: int = 3,
+    batch_size: int = 16,
+    lr: float = 1e-4,
+    n_train_samples: int = 10000,
+    device: str = 'cpu',
+    eval_every: int = 500
+):
+    """
+    Train the circuit interface layers.
+
+    Only trains:
+    - bit_extractor (embedding -> bits)
+    - bit_injector (bits -> embedding)
+    - router (circuit vs MLP weighting)
+    - op_selector (which operation)
+    """
+    print("\n" + "=" * 70)
+    print(" TRAINING CIRCUIT INTERFACE")
+    print("=" * 70)
+
+    # Freeze everything except interface layers
+    interface_params = []
+    frozen_count = 0
+    trainable_count = 0
+
+    for name, param in model.named_parameters():
+        if any(x in name for x in ['bit_extractor', 'bit_injector', 'router', 'op_selector']):
+            param.requires_grad = True
+            interface_params.append(param)
+            trainable_count += param.numel()
+        else:
+            param.requires_grad = False
+            frozen_count += param.numel()
+
+    print(f"\n  Frozen parameters: {frozen_count:,}")
+    print(f"  Trainable parameters: {trainable_count:,}")
+    print(f"  Training {len(interface_params)} parameter groups")
+
+    # Create dataset
+    print(f"\n  Creating dataset ({n_train_samples} examples)...")
+    dataset = ArithmeticDataset(tokenizer, n_samples=n_train_samples)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=collate_fn
+    )
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(interface_params, lr=lr)
+
+    # Training
+    model.to(device)
+    model.train()
+
+    global_step = 0
+    total_loss = 0
+
+    for epoch in range(n_epochs):
+        print(f"\n  Epoch {epoch + 1}/{n_epochs}")
+        print("  " + "-" * 60)
+
+        epoch_loss = 0
+        epoch_steps = 0
+
+        pbar = tqdm(dataloader, desc=f"  Training", leave=False)
+
+        for batch in pbar:
+            input_ids = batch['input_ids'].to(device)
+            labels = batch['labels'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+
+            # Forward
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels
+            )
+
+            loss = outputs.loss
+
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # Logging
+            epoch_loss += loss.item()
+            epoch_steps += 1
+            global_step += 1
+            total_loss += loss.item()
+
+            pbar.set_postfix({'loss': f'{loss.item():.4f}'})
+
+            # Periodic evaluation
+            if global_step % eval_every == 0:
+                model.eval()
+                eval_results = evaluate_arithmetic(model, tokenizer, n_problems=50, device=device)
+                print(f"\n    Step {global_step}: Loss={total_loss/eval_every:.4f}, "
+                      f"Accuracy={eval_results['accuracy']*100:.1f}%")
+                total_loss = 0
+                model.train()
+
+        avg_loss = epoch_loss / epoch_steps
+        print(f"\n  Epoch {epoch + 1} complete. Avg loss: {avg_loss:.4f}")
+
+        # End of epoch evaluation
+        model.eval()
+        eval_results = evaluate_arithmetic(model, tokenizer, n_problems=100, device=device)
+        print(f"  Evaluation: {eval_results['accuracy']*100:.1f}% "
+              f"({eval_results['correct']}/{eval_results['total']})")
+
+        if eval_results['errors']:
+            print(f"  Sample errors:")
+            for a, b, exp, got in eval_results['errors'][:3]:
+                print(f"    {a} + {b} = {exp}, model said {got}")
+
+        model.train()
+
+    print("\n" + "=" * 70)
+    print(" TRAINING COMPLETE")
+    print("=" * 70)
+
+    return model
+
+
+# =============================================================================
+# MAIN
+# =============================================================================
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Train Circuit Interface')
+    parser.add_argument('--circuit-path', type=str,
+                        default='./neural_computer.safetensors',
+                        help='Path to circuit weights')
+    parser.add_argument('--device', type=str, default='cpu',
+                        help='Device (cpu or cuda)')
+    parser.add_argument('--epochs', type=int, default=3,
+                        help='Number of epochs')
+    parser.add_argument('--batch-size', type=int, default=8,
+                        help='Batch size')
+    parser.add_argument('--lr', type=float, default=1e-4,
+                        help='Learning rate')
+    parser.add_argument('--n-samples', type=int, default=5000,
+                        help='Number of training samples')
+    args = parser.parse_args()
+
+    print("=" * 70)
+    print(" CIRCUIT-AUGMENTED LLM TRAINING")
+    print("=" * 70)
+
+    # Load model
+    print("\n[1] Loading SmolLM2-360M...")
+    model_id = "HuggingFaceTB/SmolLM2-360M"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+
+    # Baseline
+    print("\n[2] Baseline evaluation...")
+    baseline = evaluate_arithmetic(model, tokenizer, n_problems=50, device=args.device)
+    print(f"    Baseline accuracy: {baseline['accuracy']*100:.1f}%")
+
+    # Augment
+    print("\n[3] Augmenting with circuits...")
+    model = augment_smollm2_with_circuits(
+        model,
+        args.circuit_path,
+        device=args.device
+    )
+
+    # Train
+    print("\n[4] Training interface layers...")
+    model = train_interface(
+        model,
+        tokenizer,
+        n_epochs=args.epochs,
+        batch_size=args.batch_size,
+        lr=args.lr,
+        n_train_samples=args.n_samples,
+        device=args.device
+    )
+
+    # Final evaluation
+    print("\n[5] Final evaluation...")
+    final = evaluate_arithmetic(model, tokenizer, n_problems=100, device=args.device)
+    print(f"    Final accuracy: {final['accuracy']*100:.1f}%")
+    print(f"    Improvement: {baseline['accuracy']*100:.1f}% -> {final['accuracy']*100:.1f}%")
+
+    # Save
+    save_path = './circuit_augmented_smollm2.pt'
+    print(f"\n[6] Saving to {save_path}...")
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'baseline_accuracy': baseline['accuracy'],
+        'final_accuracy': final['accuracy']
+    }, save_path)
+
+    print("\nDone!")