Clean up: remove duplicates, pycache, and misplaced files

__pycache__/iron_eval.cpython-311.pyc

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

__pycache__/iron_eval.cpython-312.pyc

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

circuit_llm.py

@@ -1,606 +0,0 @@
-"""
-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)

guide.md

@@ -1,615 +0,0 @@
-# Embedding Threshold Logic Circuits into Transformer MLPs
-
-## Technical Implementation Guide
-
----
-
-## 1. Core Thesis
-
-Standard LLMs fail at arithmetic because they're interpolators—they approximate functions over training distributions rather than compute exact results. A 360M parameter model trained on internet text has seen "127 + 128 = 255" zero or few times, so it guesses "140" based on pattern matching.
-
-We solve this by embedding **frozen, proven-correct arithmetic circuits** directly into the transformer's MLP layers. The circuits use threshold logic (weighted sums + step activation), which is structurally compatible with neural network layers. We train only the **interface layers** that learn to:
-
-1. Extract operands from token embeddings
-2. Route computation through the circuits
-3. Inject results back into the residual stream
-
-The model learns **call dispatch**, not arithmetic. The arithmetic is already solved.
-
----
-
-## 2. Threshold Logic Fundamentals
-
-### 2.1 Single Threshold Gate
-
-A threshold gate computes:
-
-```
-output = 1  if  (Σ wᵢxᵢ + b) ≥ 0
-         0  otherwise
-```
-
-This is a neuron with Heaviside step activation. With integer weights `w` and bias `b`, it computes a Boolean function of binary inputs.
-
-**Example: AND gate**
-```
-w = [1, 1], b = -2
-AND(0,0) = H(0 + 0 - 2) = H(-2) = 0
-AND(0,1) = H(0 + 1 - 2) = H(-1) = 0
-AND(1,0) = H(1 + 0 - 2) = H(-1) = 0
-AND(1,1) = H(1 + 1 - 2) = H(0)  = 1
-```
-
-**Example: OR gate**
-```
-w = [1, 1], b = -1
-OR(0,0) = H(0 + 0 - 1) = H(-1) = 0
-OR(0,1) = H(0 + 1 - 1) = H(0)  = 1
-OR(1,0) = H(1 + 0 - 1) = H(0)  = 1
-OR(1,1) = H(1 + 1 - 1) = H(1)  = 1
-```
-
-### 2.2 Multi-Layer Circuits
-
-XOR is not linearly separable—it requires two layers:
-
-```
-Layer 1:
-  neuron1 (OR):   w=[1,1], b=-1   → fires if a OR b
-  neuron2 (NAND): w=[-1,-1], b=1  → fires if NOT(a AND b)
-
-Layer 2:
-  neuron3 (AND): w=[1,1], b=-2   → fires if both layer1 outputs are 1
-
-XOR(a,b) = AND(OR(a,b), NAND(a,b))
-```
-
-### 2.3 Full Adder
-
-A full adder computes `sum` and `carry_out` from inputs `a`, `b`, `carry_in`:
-
-```
-sum = a XOR b XOR cin
-cout = (a AND b) OR (cin AND (a XOR b))
-```
-
-Implementation uses two half-adders chained:
-
-```
-HA1: (a, b) → (sum1 = a XOR b, carry1 = a AND b)
-HA2: (sum1, cin) → (sum2 = sum1 XOR cin, carry2 = sum1 AND cin)
-cout = carry1 OR carry2
-final_sum = sum2
-```
-
-Each XOR is 2 layers, each AND/OR is 1 layer. Total depth: ~4 layers per full adder.
-
-### 2.4 8-bit Ripple Carry Adder
-
-Chain 8 full adders, propagating carry:
-
-```
-FA0: (a[0], b[0], 0)      → (sum[0], c0)
-FA1: (a[1], b[1], c0)     → (sum[1], c1)
-FA2: (a[2], b[2], c1)     → (sum[2], c2)
-...
-FA7: (a[7], b[7], c6)     → (sum[7], c7)
-```
-
-Total circuit depth: ~32 threshold layers (8 FAs × 4 layers each).
-
----
-
-## 3. Circuit Inventory
-
-The `neural_computer.safetensors` contains 3,122 tensors / 5,648 parameters implementing:
-
-| Category | Circuits | Tensors |
-|----------|----------|---------|
-| Boolean | AND, OR, NOT, NAND, NOR, XOR, XNOR, IMPLIES, BIIMPLIES | ~30 |
-| Arithmetic | Half adder, Full adder, Ripple carry 2/4/8-bit, 8×8 multiplier | ~800 |
-| Comparators | GT, LT, GEQ, LEQ, EQ (8-bit) | ~50 |
-| ALU | 16-operation ALU, opcode decoder, flag computation | ~400 |
-| Control | JMP, JZ, JNZ, JC, JNC, JN, JP, CALL, RET, PUSH, POP | ~200 |
-| Modular | Divisibility by 2-12 | ~600 |
-| Error Detection | Parity, CRC, Hamming, checksum | ~200 |
-| Pattern | Popcount, leading zeros, symmetry | ~150 |
-| Threshold | k-of-n gates, majority, minority | ~100 |
-
-All weights are integers. All activations are Heaviside. Verified with 6,590 exhaustive tests.
-
----
-
-## 4. Transformer Integration Architecture
-
-### 4.1 Target: SmolLM2-360M
-
-```
-Architecture: LlamaForCausalLM
-Hidden dim:   960
-Layers:       32
-Heads:        15
-MLP expansion: 4x (intermediate = 3840)
-Vocab:        49152
-Parameters:   361,821,120
-```
-
-Standard MLP block:
-```python
-def forward(x):  # x: [batch, seq, 960]
-    gate = self.gate_proj(x)      # [batch, seq, 3840]
-    up = self.up_proj(x)          # [batch, seq, 3840]
-    hidden = silu(gate) * up      # SwiGLU activation
-    return self.down_proj(hidden) # [batch, seq, 960]
-```
-
-### 4.2 Augmented MLP Block
-
-```python
-def forward(x):  # x: [batch, seq, 960]
-    # Original MLP path (unchanged)
-    mlp_out = self.down_proj(silu(self.gate_proj(x)) * self.up_proj(x))
-
-    # Circuit path (new)
-    a_bits, b_bits = self.bit_extractor(x)       # [batch, seq, 8] each
-    result_bits, carry = self.circuits.add_8bit(a_bits, b_bits)
-    flags = self.compute_flags(result_bits, carry)
-    circuit_delta = self.bit_injector(result_bits, flags)
-
-    # Routing
-    route_weights = self.router(x)  # [batch, seq, 2] softmax
-
-    # Combine
-    return mlp_out + route_weights[..., 1:2] * circuit_delta
-```
-
-### 4.3 Layer Selection
-
-We augment the **middle third** of layers (10-20 of 32):
-
-- Early layers (0-9): Token/position encoding, not arithmetic-relevant
-- Middle layers (10-20): Abstract reasoning, computation
-- Late layers (21-31): Output formatting, vocabulary projection
-
-Rationale: Arithmetic computation happens in middle layers where the model processes relationships between tokens. Early layers haven't built sufficient representations; late layers are committed to output tokens.
-
----
-
-## 5. Interface Layers (Trainable)
-
-### 5.1 BitExtractor
-
-Maps token embedding → two 8-bit operands.
-
-```python
-class BitExtractor(nn.Module):
-    def __init__(self, d_model=960):
-        self.proj = nn.Linear(d_model, 16)  # 960 → 16
-
-    def forward(self, x):
-        logits = self.proj(x)           # [batch, seq, 16]
-        bits = heaviside(logits)        # binarize with STE
-        a_bits = bits[..., :8]          # first operand
-        b_bits = bits[..., 8:]          # second operand
-        return a_bits, b_bits           # both [batch, seq, 8], LSB first
-```
-
-**What it learns**: Which embedding dimensions encode numeric magnitude. For token "127", it must learn that certain activation patterns correspond to bits `[1,1,1,1,1,1,1,0]`.
-
-**Parameters**: 960 × 16 + 16 = 15,376
-
-### 5.2 BitInjector
-
-Maps circuit outputs → embedding delta.
-
-```python
-class BitInjector(nn.Module):
-    def __init__(self, d_model=960):
-        self.proj = nn.Linear(16, d_model)  # 16 → 960
-        self.scale = nn.Parameter(torch.tensor(0.1))
-
-    def forward(self, result_bits, flags):
-        combined = torch.cat([result_bits, flags], dim=-1)  # [batch, seq, 16]
-        return self.proj(combined) * self.scale              # [batch, seq, 960]
-```
-
-**What it learns**: How to inject the result bits back into embedding space such that subsequent layers (and the final vocabulary projection) produce the correct output tokens.
-
-**Parameters**: 16 × 960 + 960 + 1 = 16,321
-
-### 5.3 Router
-
-Decides when to use circuit path.
-
-```python
-class Router(nn.Module):
-    def __init__(self, d_model=960):
-        self.net = nn.Sequential(
-            nn.Linear(d_model, 64),
-            nn.ReLU(),
-            nn.Linear(64, 2),
-            nn.Softmax(dim=-1)
-        )
-
-    def forward(self, x):
-        return self.net(x)  # [batch, seq, 2]: [mlp_weight, circuit_weight]
-```
-
-**What it learns**: "This position contains arithmetic" → route through circuits. "This is prose" → use normal MLP.
-
-**Parameters**: 960 × 64 + 64 + 64 × 2 + 2 = 61,698
-
-### 5.4 Total Trainable Parameters
-
-Per augmented layer:
-```
-BitExtractor:  15,376
-BitInjector:   16,321
-Router:        61,698
-OpSelector:    ~31,000
-───────────────────────
-Total:         ~124,395 per layer
-```
-
-For 11 augmented layers: **~1.37M trainable parameters**
-
-This is 0.38% of the model. The other 99.62% (including all circuit weights) is frozen.
-
----
-
-## 6. Gradient Flow Through Heaviside
-
-### 6.1 The Problem
-
-Heaviside has zero gradient almost everywhere:
-
-```
-H(x) = 1 if x ≥ 0 else 0
-dH/dx = 0 for x ≠ 0, undefined at x = 0
-```
-
-Standard backprop would give zero gradients to BitExtractor.
-
-### 6.2 Straight-Through Estimator (STE)
-
-We use STE: forward pass uses true Heaviside, backward pass pretends it's identity.
-
-```python
-class HeavisideSTE(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x):
-        return (x >= 0).float()  # true step function
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return grad_output  # pass gradient through unchanged
-```
-
-**Intuition**: "If making the input larger would have helped the output, increase the input." The gradient tells us the direction even though the function is flat.
-
-### 6.3 Alternative: Sigmoid Annealing
-
-During training, use sigmoid with increasing temperature:
-
-```python
-def soft_heaviside(x, temperature):
-    return torch.sigmoid(x * temperature)
-
-# temperature: 1 → 10 → 100 over training
-# At high temperature, sigmoid ≈ step function
-```
-
-This provides smoother gradients early in training, then sharpens to true binary at inference.
-
----
-
-## 7. Training Strategy
-
-### 7.1 Data Generation
-
-Generate arithmetic problems exhaustively:
-
-```python
-def generate_batch(batch_size):
-    a = torch.randint(0, 256, (batch_size,))
-    b = torch.randint(0, 256, (batch_size,))
-    result = (a + b) % 256
-
-    prompts = [f"{a[i]} + {b[i]} =" for i in range(batch_size)]
-    targets = [f" {result[i]}" for i in range(batch_size)]
-
-    return prompts, targets
-```
-
-For 8-bit addition, there are 256 × 256 = 65,536 unique problems. We can cover the entire space.
-
-### 7.2 Loss Function
-
-Standard cross-entropy on next-token prediction:
-
-```python
-outputs = model(input_ids, attention_mask=mask, labels=labels)
-loss = outputs.loss  # CE loss, only on target tokens
-```
-
-Labels are masked for prompt tokens (`-100`), so loss only backprops through the answer.
-
-### 7.3 Optimizer Configuration
-
-```python
-# Only train interface layers
-interface_params = [p for n, p in model.named_parameters()
-                    if any(x in n for x in ['bit_extractor', 'bit_injector', 'router'])]
-
-optimizer = AdamW(interface_params, lr=1e-4, weight_decay=0.01)
-scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs)
-```
-
-### 7.4 Curriculum Learning
-
-Start simple, increase difficulty:
-
-```
-Phase 1 (epochs 1-2):   Single-digit addition (0-9 + 0-9)
-Phase 2 (epochs 3-4):   Two-digit addition (0-99 + 0-99)
-Phase 3 (epochs 5-7):   Full 8-bit addition (0-255 + 0-255)
-Phase 4 (epochs 8-10):  Adversarial cases (carry chains: 127+128, 255+1)
-```
-
-This helps the interface layers learn the basic extraction pattern before tackling hard cases.
-
-### 7.5 Training Hyperparameters
-
-```
-Model:          SmolLM2-360M
-Augmented:      Layers 10-20 (11 layers)
-Trainable:      1.37M parameters
-Frozen:         362M parameters (including 5.6K circuit params)
-
-Batch size:     32
-Learning rate:  1e-4
-Epochs:         10
-Samples:        10,000 per epoch
-Warmup:         500 steps
-Device:         RTX 6000 Ada (48GB)
-
-Expected time:  ~30 minutes total
-```
-
----
-
-## 8. Forward Pass Walkthrough
-
-Input: `"127 + 128 ="`
-
-### 8.1 Tokenization
-
-```
-Tokens: ["127", " +", " 128", " ="]
-IDs:    [12700, 489, 13824, 284]  # hypothetical
-```
-
-### 8.2 Embedding
-
-```
-embeddings = embed(input_ids)  # [1, 4, 960]
-```
-
-### 8.3 Layers 0-9 (Unchanged)
-
-Standard attention + MLP, building representations.
-
-### 8.4 Layer 10 (Augmented)
-
-```python
-# After attention
-x = layer_norm(attn_output + residual)  # [1, 4, 960]
-
-# MLP path
-mlp_out = down_proj(silu(gate_proj(x)) * up_proj(x))
-
-# Circuit path
-a_bits, b_bits = bit_extractor(x)
-# Position 0 ("127"): a_bits ≈ [1,1,1,1,1,1,1,0] if well-trained
-# Position 2 ("128"): b_bits ≈ [0,0,0,0,0,0,0,1]
-# (In practice, extraction happens per-position; aggregation is learned)
-
-result_bits, carry = circuits.add_8bit(a_bits, b_bits)
-# result_bits = [1,1,1,1,1,1,1,1] = 255
-
-flags = compute_flags(result_bits, carry)
-# zero=0, negative=1, carry=1
-
-circuit_delta = bit_injector(result_bits, flags)  # [1, 4, 960]
-
-# Routing
-route = router(x)  # [1, 4, 2]
-# Position 3 ("="): route ≈ [0.1, 0.9] → use circuits
-# Position 1 ("+"): route ≈ [0.8, 0.2] → mostly MLP
-
-# Combine
-output = mlp_out + route[..., 1:2] * circuit_delta
-```
-
-### 8.5 Layers 11-31
-
-Continue processing, eventually projecting to vocabulary.
-
-### 8.6 Output
-
-```
-logits = lm_head(final_hidden)  # [1, 4, 49152]
-next_token = argmax(logits[0, 3, :])  # token after "="
-# Should decode to "255" (possibly as " 255" or "255")
-```
-
----
-
-## 9. Inference Characteristics
-
-### 9.1 Exactness
-
-At inference, Heaviside is true step function—no approximation. If BitExtractor correctly maps "127" → bits and "128" → bits, the circuit **will** output 255. The only failure mode is incorrect extraction.
-
-### 9.2 Latency
-
-Circuit computation adds ~5-10% overhead:
-- BitExtractor: 1 linear layer (960→16)
-- Circuits: ~32 threshold layers, but sparse and tiny
-- BitInjector: 1 linear layer (16→960)
-- Router: 2 linear layers
-
-The circuits have only 5,648 parameters total—negligible versus the 361M in the base model.
-
-### 9.3 Generalization
-
-Once the interface learns the mapping, it generalizes to **all** 65,536 8-bit additions. There's no memorization—the circuits compute.
-
----
-
-## 10. Evaluation Metrics
-
-### 10.1 Arithmetic Accuracy
-
-```python
-def eval_accuracy(model, n_problems=1000):
-    correct = 0
-    for _ in range(n_problems):
-        a, b = random 8-bit values
-        expected = (a + b) % 256
-        predicted = model.generate(f"{a} + {b} =")
-        if parse_int(predicted) == expected:
-            correct += 1
-    return correct / n_problems
-```
-
-**Baseline SmolLM2**: ~5-10% (guessing based on patterns)
-**Target**: >95% (circuit-accurate)
-
-### 10.2 Edge Case Performance
-
-Specifically test:
-- Carry propagation: 127+128, 255+1, 128+128
-- Zeros: 0+0, 0+255
-- Identity: x+0 for various x
-- Commutativity: verify a+b == b+a
-
-### 10.3 Non-Arithmetic Preservation
-
-Verify general capability isn't degraded:
-- Perplexity on held-out text
-- Common benchmarks (HellaSwag, etc.)
-
-The augmentation should be **additive**—circuits help arithmetic, MLP handles everything else via routing.
-
----
-
-## 11. Extension Roadmap
-
-### 11.1 Additional Operations
-
-The circuit inventory includes:
-- Subtraction (via two's complement)
-- Multiplication (8×8 → 16-bit)
-- Division (iterative subtraction)
-- Bitwise ops (AND, OR, XOR, shifts)
-- Comparisons (GT, LT, EQ)
-
-Each needs its own extraction/injection interface, or a unified interface with operation selection.
-
-### 11.2 Multi-Operand Expressions
-
-For "15 + 27 + 33 =", need:
-- Operand count detection
-- Sequential circuit invocation
-- Accumulator pattern
-
-### 11.3 Larger Bit Widths
-
-16-bit and 32-bit arithmetic require:
-- Larger circuits (or chained 8-bit)
-- Wider BitExtractor (32 or 64 output dims)
-- More training data
-
-### 11.4 Symbolic Integration
-
-Ultimate goal: the model recognizes when it needs to compute, invokes circuits, and integrates results into coherent natural language output.
-
-```
-User: "If I have 127 apples and buy 128 more, how many do I have?"
-Model: [extracts 127, 128] [routes to circuit] [gets 255]
-       "You would have 255 apples."
-```
-
----
-
-## 12. File Structure
-
-```
-8bit-threshold-computer/
-├── neural_computer.safetensors    # Frozen circuits (3,122 tensors)
-├── circuit_llm.py                 # Integration architecture
-├── train_circuit_interface.py     # Training loop
-├── iron_eval.py                   # Circuit verification (6,590 tests)
-├── skeptic_test.py                # Algebraic identity tests (127 tests)
-├── prune_weights.py               # Weight optimization
-├── tensors.txt                    # Tensor manifest
-├── guide.md                       # This document
-└── README.md                      # Project overview
-```
-
----
-
-## 13. Key Equations
-
-### Heaviside Step
-```
-H(x) = 1 if x ≥ 0 else 0
-```
-
-### Threshold Gate
-```
-f(x₁,...,xₙ) = H(Σᵢ wᵢxᵢ + b)
-```
-
-### Full Adder
-```
-sum = a ⊕ b ⊕ cᵢₙ
-cₒᵤₜ = (a ∧ b) ∨ (cᵢₙ ∧ (a ⊕ b))
-```
-
-### STE Gradient
-```
-Forward:  y = H(x)
-Backward: ∂L/∂x = ∂L/∂y
-```
-
-### Router Combination
-```
-output = mlp_out + softmax(router(x))[1] × circuit_delta
-```
-
----
-
-## 14. References
-
-1. McCulloch & Pitts (1943). "A Logical Calculus of Ideas Immanent in Nervous Activity"
-2. Muroga (1971). "Threshold Logic and Its Applications"
-3. Siegelmann & Sontag (1995). "On the Computational Power of Neural Nets"
-4. Bengio et al. (2013). "Estimating or Propagating Gradients Through Stochastic Neurons"
-5. Ma et al. (2024). "The Era of 1-bit LLMs" (BitNet b1.58)
-6. HuggingFace (2024). "SmolLM2: Small Language Models"
-
----
-
-## 15. Summary
-
-We embed a proven-correct 8-bit threshold logic computer into SmolLM2's MLP layers. The circuits are frozen; we train only the interface layers that learn call dispatch. This gives the LLM exact arithmetic capability without training it to "do math"—the math is already done.
-
-The approach is:
-- **Sound**: Circuits verified with 6,590 tests
-- **Efficient**: 1.37M trainable params, 5.6K circuit params
-- **Exact**: Heaviside at inference means no approximation error
-- **Composable**: Add more circuits (multiply, compare, etc.) with same pattern
-
-The model learns when to call the calculator, not how to calculate.

llm/todo.md

@@ -1,46 +0,0 @@
- Step 1: Install dependencies
-
-     pip install torch transformers safetensors tqdm
-
-     Step 2: Test circuit execution standalone
-
-     cd C:\Users\cnort\8bit-threshold-computer\llm
-     python circuit_llm.py --circuit-path ../neural_computer.safetensors --device cpu
-
-     Expected output:
-     - Load SmolLM2-360M
-     - Baseline arithmetic evaluation (~1-10% accuracy)
-     - Augment layers 10-20 with threshold circuits
-     - Test circuit execution (127+128=255, etc.)
-
-     Step 3: Train interface layers (CPU)
-
-     python train_circuit_interface.py \
-       --circuit-path ../neural_computer.safetensors \
-       --device cpu \
-       --epochs 3 \
-       --batch-size 4 \
-       --n-samples 2000
-
-     Expected:
-     - Freeze 362M params
-     - Train ~1.37M interface params (BitExtractor, BitInjector, Router)
-     - 3 epochs on arithmetic examples (reduced samples for CPU)
-     - Accuracy should improve from ~10% to >90%
-     - Training time: ~1-2 hours on CPU
-
-     Step 4: Save trained model
-
-     Output: circuit_augmented_smollm2.pt
-
-     Hardware Requirements
-
-     - GPU recommended (RTX 6000 Ada per guide.md)
-     - CPU possible but slow (~10x longer)
-     - ~4GB VRAM for SmolLM2-360M
-
-     Success Criteria
-
-     - Baseline accuracy: ~1-10% (SmolLM2 guessing)
-     - Post-training accuracy: >98% (circuits computing)
-     - Circuit tests pass: 127+128=255, 255+1=0, etc.

prune_weights.py

@@ -1,481 +0,0 @@
-"""
-BATCHED WEIGHT PRUNING (GPU-optimized)
-======================================
-Phase 1: Batch eval all candidates in parallel
-Phase 2: Apply all successes at once, binary search if conflicts
-"""
-
-import torch
-import time
-import argparse
-from safetensors.torch import save_file
-from iron_eval import BatchedFitnessEvaluator, create_population, load_model
-
-torch.manual_seed(0)
-
-
-def format_time(seconds):
-    if seconds < 60:
-        return f"{seconds:.1f}s"
-    elif seconds < 3600:
-        return f"{seconds/60:.1f}m"
-    else:
-        return f"{seconds/3600:.1f}h"
-
-
-def format_eta(elapsed, done, total):
-    if done == 0:
-        return "calculating..."
-    rate = done / elapsed
-    remaining = (total - done) / rate
-    return format_time(remaining)
-
-
-def apply_reductions(model, reductions):
-    """Apply a list of (name, flat_idx, shape, old_val) reductions."""
-    for name, flat_idx, shape, old_val in reductions:
-        new_val = old_val - 1 if old_val > 0 else old_val + 1
-        flat = model[name].flatten()
-        if flat[flat_idx].item() == old_val:
-            flat[flat_idx] = new_val
-            model[name] = flat.view(shape)
-
-
-def revert_reductions(model, reductions):
-    """Revert a list of reductions."""
-    for name, flat_idx, shape, old_val in reductions:
-        flat = model[name].flatten()
-        new_val = old_val - 1 if old_val > 0 else old_val + 1
-        if flat[flat_idx].item() == new_val:
-            flat[flat_idx] = old_val
-            model[name] = flat.view(shape)
-
-
-def check_fitness(model, evaluator, device):
-    """Check model fitness."""
-    torch.manual_seed(0)
-    pop = create_population(model, 1, device)
-    return evaluator.evaluate(pop, debug=False)[0].item()
-
-
-def sequential_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
-    """
-    Sequential fallback - tests and applies reductions one at a time.
-    Slower but guarantees no interaction bugs.
-    """
-    accepted = []
-    for i, (name, flat_idx, shape, old_val) in enumerate(candidates):
-        apply_reductions(model, [(name, flat_idx, shape, old_val)])
-        fitness = check_fitness(model, evaluator, device)
-        if fitness >= 0.9999:
-            accepted.append((name, flat_idx, shape, old_val))
-            if (i + 1) % 50 == 0:
-                current_mag = sum(t.abs().sum().item() for t in model.values())
-                reduction_pct = 100 * (1 - current_mag / base_magnitude)
-                print(f"          Sequential: {len(accepted)}/{i+1} accepted | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
-        else:
-            revert_reductions(model, [(name, flat_idx, shape, old_val)])
-    return accepted
-
-
-def batched_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
-    """
-    Batched binary search - evaluates multiple branches in parallel.
-    Uses BFS instead of DFS to maximize batching opportunities.
-    Verifies cumulative effect after each batch to prevent interaction bugs.
-    """
-    if len(candidates) == 0:
-        return []
-
-    # First try all at once
-    print(f"        Trying {len(candidates)} reductions at once...")
-    apply_reductions(model, candidates)
-    fitness = check_fitness(model, evaluator, device)
-
-    if fitness >= 0.9999:
-        current_mag = sum(t.abs().sum().item() for t in model.values())
-        reduction_pct = 100 * (1 - current_mag / base_magnitude)
-        print(f"        ALL {len(candidates)} OK | fitness={fitness:.6f} | "
-              f"mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
-        return candidates
-
-    # Conflict - revert and use batched BFS
-    revert_reductions(model, candidates)
-    print(f"        CONFLICT (fitness={fitness:.6f}), starting batched resolution...")
-
-    accepted = []
-    # Queue of (candidate_list, depth) to process
-    pending = [(candidates, 0)]
-
-    while pending:
-        # Collect all pending groups for batch evaluation
-        to_eval = []
-        for group, depth in pending:
-            if len(group) == 0:
-                continue
-            elif len(group) == 1:
-                to_eval.append((group, depth, 'single'))
-            else:
-                to_eval.append((group, depth, 'group'))
-
-        pending = []
-
-        if not to_eval:
-            break
-
-        # Build batch: create model variants for each group
-        batch_size = len(to_eval)
-        print(f"          Batch evaluating {batch_size} groups...")
-
-        # Create population for batch eval
-        pop = {}
-        for name, tensor in model.items():
-            pop[name] = tensor.unsqueeze(0).expand(batch_size, *tensor.shape).clone().to(device)
-
-        # Apply each group's reductions to its population slot
-        for idx, (group, depth, gtype) in enumerate(to_eval):
-            for name, flat_idx, shape, old_val in group:
-                new_val = old_val - 1 if old_val > 0 else old_val + 1
-                flat_view = pop[name][idx].flatten()
-                # Check if not already modified in base model
-                base_val = model[name].flatten()[flat_idx].item()
-                if base_val == old_val:
-                    flat_view[flat_idx] = new_val
-
-        # Batch evaluate
-        torch.manual_seed(0)
-        fitnesses = evaluator.evaluate(pop, debug=False)
-
-        # Process results - collect accepted groups first, then verify
-        batch_accepted = []
-        ok_count = 0
-        conflict_count = 0
-        fail_count = 0
-
-        for idx, (group, depth, gtype) in enumerate(to_eval):
-            fit = fitnesses[idx].item()
-            indent = "          " + "  " * depth
-
-            if fit >= 0.9999:
-                batch_accepted.append((group, depth, indent))
-                ok_count += len(group)
-            else:
-                if len(group) == 1:
-                    name, flat_idx, shape, old_val = group[0]
-                    print(f"{indent}[1/1] FAIL {name}[{flat_idx}] | fitness={fit:.6f}")
-                    fail_count += 1
-                else:
-                    mid = len(group) // 2
-                    left = group[:mid]
-                    right = group[mid:]
-                    print(f"{indent}CONFLICT ({len(group)}) fitness={fit:.6f} -> split {len(left)}+{len(right)}")
-                    pending.append((left, depth + 1))
-                    pending.append((right, depth + 1))
-                    conflict_count += 1
-
-        # Apply all batch-accepted reductions
-        all_batch_reductions = []
-        for group, depth, indent in batch_accepted:
-            apply_reductions(model, group)
-            all_batch_reductions.extend(group)
-
-        # Verify cumulative effect
-        if all_batch_reductions:
-            verify_fitness = check_fitness(model, evaluator, device)
-            if verify_fitness >= 0.9999:
-                # All good - commit these reductions
-                for group, depth, indent in batch_accepted:
-                    current_mag = sum(t.abs().sum().item() for t in model.values())
-                    reduction_pct = 100 * (1 - current_mag / base_magnitude)
-                    if len(group) == 1:
-                        name, flat_idx, shape, old_val = group[0]
-                        print(f"{indent}[1/1] OK {name}[{flat_idx}] | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
-                    else:
-                        print(f"{indent}ALL {len(group)} OK | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
-                accepted.extend(all_batch_reductions)
-                print(f"          Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
-            else:
-                # Interaction bug detected - revert and use sequential fallback
-                print(f"          INTERACTION BUG detected (batch fitness={verify_fitness:.6f})")
-                print(f"          Reverting {len(all_batch_reductions)} reductions, falling back to sequential...")
-                revert_reductions(model, all_batch_reductions)
-
-                # Process each group sequentially
-                seq_accepted = sequential_conflict_resolution(
-                    model, evaluator, device, all_batch_reductions, base_magnitude
-                )
-                accepted.extend(seq_accepted)
-                print(f"          Sequential fallback: {len(seq_accepted)}/{len(all_batch_reductions)} accepted")
-        else:
-            print(f"          Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
-
-    return accepted
-
-
-def prune_weights(
-    passes: int = 10,
-    batch_size: int = 5000,
-    device: str = 'cuda',
-    checkpoint_path: str = "D:/8bit-threshold-computer/pruned.safetensors"
-):
-    print("=" * 80)
-    print(" BATCHED WEIGHT PRUNING (GPU-optimized)")
-    print("=" * 80)
-    print(f" Device: {device}")
-    print(f" Batch size: {batch_size}")
-    print(f" Max passes: {passes}")
-    print("=" * 80)
-
-    # Load model
-    print("\n[1/4] LOADING MODEL...")
-    load_start = time.perf_counter()
-    model = load_model()
-    load_time = time.perf_counter() - load_start
-
-    n_params = sum(t.numel() for t in model.values())
-    n_tensors = len(model)
-    base_magnitude = sum(t.abs().sum().item() for t in model.values())
-    base_max = max(t.abs().max().item() for t in model.values())
-    nonzero_params = sum((t != 0).sum().item() for t in model.values())
-
-    print(f"       Loaded in {load_time:.2f}s")
-    print(f"       Tensors: {n_tensors}")
-    print(f"       Parameters: {n_params}")
-    print(f"       Non-zero parameters: {nonzero_params}")
-    print(f"       Total magnitude: {base_magnitude:.0f}")
-    print(f"       Max weight: {base_max:.0f}")
-
-    # Initialize evaluator
-    print("\n[2/4] INITIALIZING EVALUATOR...")
-    eval_start = time.perf_counter()
-    evaluator = BatchedFitnessEvaluator(device=device)
-    eval_time = time.perf_counter() - eval_start
-    print(f"       Initialized in {eval_time:.2f}s")
-
-    # Verify initial fitness
-    print("\n[3/4] VERIFYING BASE MODEL...")
-    initial_fitness = check_fitness(model, evaluator, device)
-    print(f"       Fitness: {initial_fitness:.6f}")
-
-    if initial_fitness < 0.9999:
-        print(f"       ERROR: Base model fitness {initial_fitness:.6f} < 0.9999")
-        return None
-
-    print(f"       STATUS: PASS")
-
-    # Build parameter list
-    print("\n[4/4] BUILDING PARAMETER INDEX...")
-    param_list = []
-    for name, tensor in model.items():
-        flat = tensor.flatten()
-        for i in range(len(flat)):
-            param_list.append((name, i, tensor.shape))
-    print(f"       Indexed {len(param_list)} parameters")
-
-    # Main pruning loop
-    print("\n" + "=" * 80)
-    print(" PRUNING STARTED")
-    print("=" * 80)
-
-    total_reductions = 0
-    pruning_start = time.perf_counter()
-
-    for pass_num in range(passes):
-        torch.manual_seed(0)
-        pass_start = time.perf_counter()
-
-        print(f"\n{'='*80}")
-        print(f" PASS {pass_num + 1}/{passes}")
-        print(f"{'='*80}")
-
-        # Count candidates
-        candidates = []
-        for name, idx, shape in param_list:
-            flat = model[name].flatten()
-            val = flat[idx].item()
-            if val != 0:
-                candidates.append((name, idx, shape, val))
-
-        n_candidates = len(candidates)
-        print(f"\n    Candidates: {n_candidates} non-zero weights")
-
-        if n_candidates == 0:
-            print(f"    No candidates remaining. Stopping.")
-            break
-
-        # Phase 1: Batch evaluation
-        print(f"\n    PHASE 1: Batch evaluation (testing each reduction independently)")
-        print(f"    " + "-" * 60)
-        phase1_start = time.perf_counter()
-        successful_candidates = []
-        n_batches = (n_candidates + batch_size - 1) // batch_size
-
-        for batch_idx, batch_start_idx in enumerate(range(0, n_candidates, batch_size)):
-            batch = candidates[batch_start_idx:batch_start_idx + batch_size]
-            batch_len = len(batch)
-            batch_start_time = time.perf_counter()
-
-            # Build population
-            pop = {}
-            for name, tensor in model.items():
-                pop[name] = tensor.unsqueeze(0).expand(batch_len, *tensor.shape).clone().to(device)
-
-            # Apply reductions
-            for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
-                new_val = old_val - 1 if old_val > 0 else old_val + 1
-                flat_view = pop[name][pop_idx].flatten()
-                flat_view[flat_idx] = new_val
-
-            # Evaluate
-            torch.manual_seed(0)
-            if device == 'cuda':
-                torch.cuda.synchronize()
-            fitness = evaluator.evaluate(pop, debug=False)
-            if device == 'cuda':
-                torch.cuda.synchronize()
-
-            # Collect successes
-            batch_successes = 0
-            for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
-                if fitness[pop_idx].item() >= 0.9999:
-                    successful_candidates.append((name, flat_idx, shape, old_val))
-                    batch_successes += 1
-
-            batch_time = time.perf_counter() - batch_start_time
-            elapsed = time.perf_counter() - phase1_start
-            done = batch_start_idx + batch_len
-            eta = format_eta(elapsed, done, n_candidates)
-            throughput = batch_len / batch_time
-
-            print(f"      Batch {batch_idx + 1}/{n_batches}: "
-                  f"{batch_successes}/{batch_len} passed ({100*batch_successes/batch_len:.1f}%) | "
-                  f"Total OK: {len(successful_candidates)} | "
-                  f"Progress: {done}/{n_candidates} ({100*done/n_candidates:.1f}%) | "
-                  f"Speed: {throughput:.0f}/s | "
-                  f"ETA: {eta}")
-
-        phase1_time = time.perf_counter() - phase1_start
-        print(f"\n    Phase 1 complete: {len(successful_candidates)}/{n_candidates} candidates "
-              f"({100*len(successful_candidates)/n_candidates:.1f}%) in {format_time(phase1_time)}")
-
-        # Phase 2: Apply with conflict resolution
-        if len(successful_candidates) == 0:
-            print(f"\n    No reductions possible. Stopping.")
-            break
-
-        print(f"\n    PHASE 2: Apply reductions with conflict resolution")
-        print(f"    " + "-" * 60)
-        phase2_start = time.perf_counter()
-
-        accepted = batched_conflict_resolution(model, evaluator, device, successful_candidates, base_magnitude)
-        pass_reductions = len(accepted)
-
-        phase2_time = time.perf_counter() - phase2_start
-        print(f"\n    Phase 2 complete: {pass_reductions} reductions applied in {format_time(phase2_time)}")
-
-        # Pass summary
-        total_reductions += pass_reductions
-        current_magnitude = sum(t.abs().sum().item() for t in model.values())
-        current_nonzero = sum((t != 0).sum().item() for t in model.values())
-        pass_time = time.perf_counter() - pass_start
-        reduction_pct = 100 * (1 - current_magnitude / base_magnitude)
-
-        print(f"\n    PASS {pass_num + 1} SUMMARY:")
-        print(f"      Reductions this pass: {pass_reductions}")
-        print(f"      Total reductions: {total_reductions}")
-        print(f"      Current magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
-        print(f"      Current non-zero: {current_nonzero}")
-        print(f"      Pass time: {format_time(pass_time)}")
-
-        # Verify after pass
-        print(f"\n    Verifying model integrity...")
-        fitness = check_fitness(model, evaluator, device)
-        print(f"      Fitness: {fitness:.6f} {'PASS' if fitness >= 0.9999 else 'FAIL'}")
-
-        # Save checkpoint after each pass
-        checkpoint_name = checkpoint_path.replace('.safetensors', f'_pass{pass_num + 1}.safetensors')
-        print(f"\n    Saving checkpoint: {checkpoint_name}")
-        save_file(model, checkpoint_name)
-        print(f"      Saved. Magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
-
-        # Also save as "latest" for easy access
-        latest_path = checkpoint_path.replace('.safetensors', '_latest.safetensors')
-        save_file(model, latest_path)
-        print(f"      Also saved as: {latest_path}")
-
-        if pass_reductions == 0:
-            print(f"\n    No reductions achieved. Stopping early.")
-            break
-
-    # Final summary
-    pruning_time = time.perf_counter() - pruning_start
-    final_magnitude = sum(t.abs().sum().item() for t in model.values())
-    final_max = max(t.abs().max().item() for t in model.values())
-    final_nonzero = sum((t != 0).sum().item() for t in model.values())
-    reduction_pct = 100 * (1 - final_magnitude / base_magnitude)
-
-    print("\n" + "=" * 80)
-    print(" PRUNING COMPLETE")
-    print("=" * 80)
-    print(f"\n RESULTS:")
-    print(f"   Original magnitude:  {base_magnitude:.0f}")
-    print(f"   Final magnitude:     {final_magnitude:.0f}")
-    print(f"   Reduction:           {reduction_pct:.2f}%")
-    print(f"   Total reductions:    {total_reductions}")
-    print(f"   Original non-zero:   {nonzero_params}")
-    print(f"   Final non-zero:      {final_nonzero}")
-    print(f"   Zeros created:       {nonzero_params - final_nonzero}")
-    print(f"   Max weight:          {final_max:.0f}")
-    print(f"   Total time:          {format_time(pruning_time)}")
-
-    # Save
-    print(f"\n SAVING to {checkpoint_path}...")
-    save_file(model, checkpoint_path)
-    print(f"   Saved.")
-
-    # Final verification
-    print(f"\n FINAL VERIFICATION...")
-    from safetensors import safe_open
-    f = safe_open(checkpoint_path, framework='numpy')
-    verify_model = {name: torch.tensor(f.get_tensor(name)).float() for name in f.keys()}
-    verify_fitness = check_fitness(verify_model, evaluator, device)
-    print(f"   Fitness: {verify_fitness:.6f}")
-
-    if verify_fitness >= 0.9999:
-        print(f"   STATUS: PASS")
-    else:
-        print(f"   STATUS: FAIL - Model corrupted!")
-
-    print("\n" + "=" * 80)
-    return model
-
-
-MAX_BATCH_SIZE = 80000
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description='Batched Weight Pruning')
-    parser.add_argument('--passes', type=int, default=10,
-                        help='Maximum pruning passes (default: 10)')
-    parser.add_argument('--batch_size', type=int, default=80000,
-                        help=f'Batch size for parallel evaluation (default: 80000, max: {MAX_BATCH_SIZE})')
-    parser.add_argument('--device', type=str, default='cuda',
-                        help='Device: cuda or cpu (default: cuda)')
-    parser.add_argument('--output', type=str,
-                        default='D:/8bit-threshold-computer/pruned.safetensors',
-                        help='Output path')
-    args = parser.parse_args()
-
-    if args.batch_size > MAX_BATCH_SIZE:
-        print(f"WARNING: batch_size {args.batch_size} exceeds maximum {MAX_BATCH_SIZE}. Clamping.")
-        args.batch_size = MAX_BATCH_SIZE
-
-    print(f"\nStarting at {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
-
-    prune_weights(
-        passes=args.passes,
-        batch_size=args.batch_size,
-        device=args.device,
-        checkpoint_path=args.output
-    )
-
-    print(f"\nFinished at {time.strftime('%Y-%m-%d %H:%M:%S')}")

stress_test.py

@@ -1,367 +0,0 @@
-"""
-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

@@ -1,306 +0,0 @@
-"""
-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!")