v2: 293/400 tasks solved (was 128). Added variable-shape conv and diff-shape conv solvers.

neurogolf_solver.py

@@ -1,12 +1,19 @@
 #!/usr/bin/env python3
 """
-ARC-AGI NeuroGolf Championship - Complete Solver
+ARC-AGI NeuroGolf Championship - Complete Solver v2
 Format: [1,10,30,30] one-hot input/output, opset 10, IR version 10.
-Pipeline: Slice -> Conv -> ArgMax -> OneHot -> Pad
+
+Solvers:
+  - Analytical: identity, constant, color_map, transpose, flip, rotate, tile, upscale, concat, spatial_gather
+  - Conv (fixed shape): Slice -> Conv -> ArgMax -> OneHot -> Pad
+  - Conv (variable shape): Conv(30x30) -> ArgMax -> OneHot -> Mul(mask)  [NEW]
+  - Conv (diff shape): Slice -> Conv -> Slice(crop) -> ArgMax -> OneHot -> Pad  [NEW]
+
+Results: 293/400 tasks solved (was 128/400 in v1)
 
 Usage:
   python neurogolf_solver.py --data_dir ARC-AGI/data/training/ --output_dir submission
-  python neurogolf_solver.py --data_dir ARC-AGI/data/training/ --output_dir submission --device cuda --conv_budget 60
+  python neurogolf_solver.py --data_dir ARC-AGI/data/training/ --output_dir submission --conv_budget 60
 """
 
 import json, os, sys, math, time, argparse
@@ -23,13 +30,9 @@ IR = 10
 OPSET = [helper.make_opsetid("", 10)]
 
 def get_providers():
-    available = ort.get_available_providers()
-    if 'CUDAExecutionProvider' in available:
-        return ['CUDAExecutionProvider', 'CPUExecutionProvider']
-    return ['CPUExecutionProvider']
+    return ['CPUExecutionProvider']  # CPU is faster for tiny 30x30 grids
 
 ORT_PROVIDERS = get_providers()
-print(f"ONNX Runtime providers: {ORT_PROVIDERS}")
 
 def load_tasks_dir(data_dir):
     files = sorted(f for f in os.listdir(data_dir) if f.endswith('.json'))
@@ -92,7 +95,7 @@ def fixed_shapes(td):
     return list(shapes)[0] if len(shapes) == 1 else None
 
 # ============================================================
-# SOLVERS
+# ANALYTICAL SOLVERS
 # ============================================================
 
 def s_identity(td):
@@ -295,48 +298,71 @@ def s_constant(td):
     return mk(nodes, inits)
 
 # ============================================================
-# CONV SOLVER
+# CONV SOLVER (fixed shape) - Slice -> Conv -> ArgMax -> OneHot -> Pad
 # ============================================================
 
-def solve_conv(td, path, time_budget=30.0, try_bias=True):
+def _lstsq_conv(exs_raw, ks, use_bias, use_full_30=False):
+    """Shared lstsq conv fitting. Returns (Wconv, B) or None."""
+    pad = ks // 2
+    feat = 10 * ks * ks + (1 if use_bias else 0)
+    if feat > 20000: return None
+    
+    patches, targets = [], []
+    for inp_g, out_g in exs_raw:
+        ih, iw = inp_g.shape
+        if use_full_30:
+            oh_full = np.zeros((10, GH, GW), dtype=np.float64)
+            for c in range(10): oh_full[c, :ih, :iw] = (inp_g == c)
+            oh_pad = np.pad(oh_full, ((0,0),(pad,pad),(pad,pad)))
+        else:
+            oh_enc = np.zeros((10, ih, iw), dtype=np.float64)
+            for c in range(10): oh_enc[c] = (inp_g == c)
+            oh_pad = np.pad(oh_enc, ((0,0),(pad,pad),(pad,pad)))
+        
+        oh, ow = out_g.shape
+        for r in range(oh):
+            for c in range(ow):
+                p = oh_pad[:, r:r+ks, c:c+ks].flatten()
+                if use_bias: p = np.append(p, 1.0)
+                patches.append(p)
+                targets.append(int(out_g[r, c]))
+    
+    n_patches = len(patches)
+    if feat > 5000 and n_patches > 2000: return None
+    
+    P = np.array(patches, dtype=np.float64)
+    T = np.array(targets, dtype=np.int64)
+    T_oh = np.zeros((len(T), 10), dtype=np.float64)
+    for i, t in enumerate(T): T_oh[i, t] = 1.0
+    
+    WT = np.linalg.lstsq(P, T_oh, rcond=None)[0]
+    if not np.array_equal(np.argmax(P @ WT, axis=1), T): return None
+    
+    if use_bias:
+        Wconv = WT[:-1].T.reshape(10, 10, ks, ks).astype(np.float32)
+        B = WT[-1].astype(np.float32)
+    else:
+        Wconv = WT.T.reshape(10, 10, ks, ks).astype(np.float32)
+        B = None
+    return Wconv, B
+
+def solve_conv_fixed(td, path, time_budget=30.0):
+    """Fixed-shape conv: Slice -> Conv -> ArgMax -> OneHot -> Pad."""
     exs = get_exs(td)
     for inp, out in exs:
         if inp.shape != out.shape: return None
     shapes = set(inp.shape for inp, _ in exs)
     if len(shapes) != 1: return None
     IH, IW = shapes.pop()
+    
     t_start = time.time()
-    for use_bias in ([False, True] if try_bias else [False]):
+    for use_bias in [False, True]:
         for ks in [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29]:
             if time.time() - t_start > time_budget: return None
+            result = _lstsq_conv(exs, ks, use_bias, use_full_30=False)
+            if result is None: continue
+            Wconv, B = result
             pad = ks // 2
-            feat = 10 * ks * ks + (1 if use_bias else 0)
-            n_grid = sum(inp.size for inp, _ in exs)
-            if feat > 20000 or (feat > 5000 and n_grid > 2000): continue
-            patches, targets = [], []
-            for inp_g, out_g in exs:
-                ih, iw = inp_g.shape
-                oh_enc = np.zeros((10, ih, iw), dtype=np.float64)
-                for c in range(10): oh_enc[c] = (inp_g == c)
-                oh_pad = np.pad(oh_enc, ((0,0),(pad,pad),(pad,pad)))
-                for r in range(ih):
-                    for c in range(iw):
-                        p = oh_pad[:, r:r+ks, c:c+ks].flatten()
-                        if use_bias: p = np.append(p, 1.0)
-                        patches.append(p)
-                        targets.append(int(out_g[r, c]))
-            P = np.array(patches, dtype=np.float64)
-            T = np.array(targets, dtype=np.int64)
-            T_oh = np.zeros((len(T), 10), dtype=np.float64)
-            for i, t in enumerate(T): T_oh[i, t] = 1.0
-            WT = np.linalg.lstsq(P, T_oh, rcond=None)[0]
-            if not np.array_equal(np.argmax(P @ WT, axis=1), T): continue
-            if use_bias:
-                Wconv = WT[:-1].T.reshape(10, 10, ks, ks).astype(np.float32)
-                B = WT[-1].astype(np.float32)
-            else:
-                Wconv = WT.T.reshape(10, 10, ks, ks).astype(np.float32)
-                B = None
             pad_h, pad_w = GH - IH, GW - IW
             inits = [
                 numpy_helper.from_array(np.array([0,0,0,0], dtype=np.int64), 'sl_st'),
@@ -361,6 +387,134 @@ def solve_conv(td, path, time_budget=30.0, try_bias=True):
             if validate(path, td): return model
     return None
 
+# ============================================================
+# CONV SOLVER (variable shape) - Conv(30x30) -> ArgMax -> OneHot -> Mul(mask)
+# ============================================================
+
+def solve_conv_variable(td, path, time_budget=30.0):
+    """Variable-shape conv: works on full 30x30 one-hot, dynamic mask from input."""
+    exs = get_exs(td)
+    for inp, out in exs:
+        if inp.shape != out.shape: return None
+    
+    t_start = time.time()
+    for use_bias in [False, True]:
+        for ks in [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29]:
+            if time.time() - t_start > time_budget: return None
+            result = _lstsq_conv(exs, ks, use_bias, use_full_30=True)
+            if result is None: continue
+            Wconv, B = result
+            pad = ks // 2
+            inits = [
+                numpy_helper.from_array(Wconv, 'W'),
+                numpy_helper.from_array(np.array(10, dtype=np.int64), 'depth'),
+                numpy_helper.from_array(np.array([0.0, 1.0], dtype=np.float32), 'ohvals'),
+            ]
+            conv_inputs = ['input', 'W']
+            if B is not None:
+                inits.append(numpy_helper.from_array(B, 'B'))
+                conv_inputs.append('B')
+            nodes = [
+                helper.make_node('ReduceSum', ['input'], ['mask'], axes=[1], keepdims=1),
+                helper.make_node('Conv', conv_inputs, ['co'], kernel_shape=[ks,ks], pads=[pad]*4),
+                helper.make_node('ArgMax', ['co'], ['am'], axis=1, keepdims=0),
+                helper.make_node('OneHot', ['am', 'depth', 'ohvals'], ['oh_out'], axis=1),
+                helper.make_node('Mul', ['oh_out', 'mask'], ['output']),
+            ]
+            model = mk(nodes, inits)
+            onnx.save(model, path)
+            if validate(path, td): return model
+    return None
+
+# ============================================================
+# CONV SOLVER (diff shape, fixed) - output smaller than input
+# ============================================================
+
+def solve_conv_diffshape(td, path, time_budget=30.0):
+    """Diff-shape conv for fixed io shapes where output is smaller."""
+    sp = fixed_shapes(td)
+    if sp is None: return None
+    (IH, IW), (OH, OW) = sp
+    if IH == OH and IW == OW: return None
+    if OH > IH or OW > IW: return None
+    if OH > 30 or OW > 30: return None
+    
+    exs = get_exs(td)
+    t_start = time.time()
+    
+    for dr_off, dc_off in [(0, 0), ((IH-OH)//2, (IW-OW)//2)]:
+        for use_bias in [False, True]:
+            for ks in [1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21]:
+                if time.time() - t_start > time_budget: return None
+                pad = ks // 2
+                feat = 10 * ks * ks + (1 if use_bias else 0)
+                if feat > 10000: continue
+                
+                patches, targets = [], []
+                valid = True
+                for inp_g, out_g in exs:
+                    oh_enc = np.zeros((10, IH, IW), dtype=np.float64)
+                    for c in range(10): oh_enc[c] = (inp_g == c)
+                    oh_pad = np.pad(oh_enc, ((0,0),(pad,pad),(pad,pad)))
+                    for r in range(OH):
+                        for c in range(OW):
+                            sr, sc = r + dr_off, c + dc_off
+                            if sr < 0 or sr >= IH or sc < 0 or sc >= IW:
+                                valid = False; break
+                            p = oh_pad[:, sr:sr+ks, sc:sc+ks].flatten()
+                            if use_bias: p = np.append(p, 1.0)
+                            patches.append(p)
+                            targets.append(int(out_g[r, c]))
+                        if not valid: break
+                    if not valid: break
+                if not valid: continue
+                
+                n_patches = len(patches)
+                if feat > 5000 and n_patches > 2000: continue
+                
+                P = np.array(patches, dtype=np.float64)
+                T = np.array(targets, dtype=np.int64)
+                T_oh = np.zeros((len(T), 10), dtype=np.float64)
+                for i, t in enumerate(T): T_oh[i, t] = 1.0
+                
+                WT = np.linalg.lstsq(P, T_oh, rcond=None)[0]
+                if not np.array_equal(np.argmax(P @ WT, axis=1), T): continue
+                
+                if use_bias:
+                    Wconv = WT[:-1].T.reshape(10, 10, ks, ks).astype(np.float32)
+                    B = WT[-1].astype(np.float32)
+                else:
+                    Wconv = WT.T.reshape(10, 10, ks, ks).astype(np.float32)
+                    B = None
+                
+                pad_h, pad_w = GH - OH, GW - OW
+                inits = [
+                    numpy_helper.from_array(np.array([0,0,0,0], dtype=np.int64), 'sl_st'),
+                    numpy_helper.from_array(np.array([1,10,IH,IW], dtype=np.int64), 'sl_en'),
+                    numpy_helper.from_array(Wconv, 'W'),
+                    numpy_helper.from_array(np.array(10, dtype=np.int64), 'depth'),
+                    numpy_helper.from_array(np.array([0.0, 1.0], dtype=np.float32), 'ohvals'),
+                    numpy_helper.from_array(np.array([0,0,dr_off,dc_off], dtype=np.int64), 'cr_st'),
+                    numpy_helper.from_array(np.array([1,10,dr_off+OH,dc_off+OW], dtype=np.int64), 'cr_en'),
+                ]
+                conv_inputs = ['grid', 'W']
+                if B is not None:
+                    inits.append(numpy_helper.from_array(B, 'B'))
+                    conv_inputs.append('B')
+                
+                nodes = [
+                    helper.make_node('Slice', ['input','sl_st','sl_en'], ['grid']),
+                    helper.make_node('Conv', conv_inputs, ['co'], kernel_shape=[ks,ks], pads=[pad]*4),
+                    helper.make_node('Slice', ['co','cr_st','cr_en'], ['co_crop']),
+                    helper.make_node('ArgMax', ['co_crop'], ['am'], axis=1, keepdims=0),
+                    helper.make_node('OneHot', ['am','depth','ohvals'], ['oh_out'], axis=1),
+                    helper.make_node('Pad', ['oh_out'], ['output'], pads=[0,0,0,0,0,0,pad_h,pad_w], value=0.0),
+                ]
+                model = mk(nodes, inits)
+                onnx.save(model, path)
+                if validate(path, td): return model
+    return None
+
 # ============================================================
 # GATHER HELPERS
 # ============================================================
@@ -431,6 +585,8 @@ ANALYTICAL_SOLVERS = [
 def solve_task(tn, td, outdir, conv_budget=30.0):
     os.makedirs(outdir, exist_ok=True)
     path = os.path.join(outdir, f"task{tn:03d}.onnx")
+    
+    # 1. Try analytical solvers (fast, tiny models)
     for sname, sfn in ANALYTICAL_SOLVERS:
         try:
             model = sfn(td)
@@ -438,8 +594,31 @@ def solve_task(tn, td, outdir, conv_budget=30.0):
             onnx.save(model, path)
             if validate(path, td): return True, sname, os.path.getsize(path)
         except: pass
-    model = solve_conv(td, path, time_budget=conv_budget)
-    if model is not None: return True, 'conv', os.path.getsize(path)
+    
+    # 2. Determine task shape category
+    exs = get_exs(td)
+    same_shape = all(inp.shape == out.shape for inp, out in exs)
+    shapes = set(inp.shape for inp, _ in exs)
+    fixed_in = len(shapes) == 1
+    
+    if same_shape:
+        if fixed_in:
+            # Fixed same-shape: use original conv (Slice->Conv->Pad)
+            model = solve_conv_fixed(td, path, time_budget=conv_budget)
+            if model is not None: return True, 'conv_fixed', os.path.getsize(path)
+        # Always try variable-shape conv for same-shape tasks
+        model = solve_conv_variable(td, path, time_budget=conv_budget)
+        if model is not None: return True, 'conv_var', os.path.getsize(path)
+    else:
+        # Different shapes
+        sp = fixed_shapes(td)
+        if sp is not None:
+            (IH,IW),(OH,OW) = sp
+            if OH <= IH and OW <= IW:
+                # Output smaller: try diff-shape conv
+                model = solve_conv_diffshape(td, path, time_budget=conv_budget)
+                if model is not None: return True, 'conv_diff', os.path.getsize(path)
+    
     return False, None, None
 
 def main():