Fix self_tile: use ConvTranspose for block-upscale (onnx_tool profiles ConvTranspose correctly)

own-solver/neurogolf_solver/solvers/wave3.py

@@ -1,16 +1,9 @@
 #!/usr/bin/env python3
 """Wave 3: Split-and-merge + Self-tile solvers.
 
-- s_split_and_merge: split grid at separator, merge halves via OR/AND/XOR
-- s_self_tile: output = input tiled by its own non-bg mask
-- s_self_tile_majority: output = input tiled by majority-color mask
+Uses ConvTranspose(stride) for block-upscaling masks (onnx_tool compatible).
 
-NOTE: Uses Reshape+Tile for block-upscaling instead of Resize (onnx_tool can't profile Resize).
-
-Validated:
-- e98196ab (OR, horizontal split) — 262/262
-- 007bbfb7 (self-tile non-bg mask) — 262/262
-- c3e719e8 (self-tile majority mask) — 262/262
+Validated: 007bbfb7 (262/262), c3e719e8 (262/262), e98196ab (262/262)
 """
 
 import numpy as np
@@ -20,28 +13,17 @@ from ..data_loader import get_exs, fixed_shapes
 from ..constants import GH, GW
 
 
-# =============================================================================
-# SELF-TILE SOLVER (non-bg mask)
-# =============================================================================
-
 def s_self_tile(td):
-    """Self-tiling: output = input tiled by its own non-bg mask.
-
-    Each non-bg pixel in input becomes a full copy of the input.
-    Each bg pixel becomes a block of zeros.
-    Output size = IH*IH x IW*IW.
-    """
+    """Self-tiling: each non-bg pixel → full copy of input. Output = IH^2 x IW^2."""
     exs = get_exs(td)
     sp = fixed_shapes(td)
     if sp is None:
         return None
     (IH, IW), (OH, OW) = sp
-
     if OH != IH * IH or OW != IW * IW:
         return None
     if OH > 30 or OW > 30:
         return None
-
     for inp, out in exs:
         expected = np.zeros((OH, OW), dtype=np.int64)
         for r in range(IH):
@@ -51,63 +33,44 @@ def s_self_tile(td):
         if not np.array_equal(expected, out):
             return None
 
-    non_bg_ch_mask = np.ones((1, 10, 1, 1), dtype=np.float32)
-    non_bg_ch_mask[0, 0, 0, 0] = 0.0
+    non_bg_ch = np.ones((1, 10, 1, 1), dtype=np.float32)
+    non_bg_ch[0, 0, 0, 0] = 0.0
+    up_k = np.ones((1, 1, IH, IW), dtype=np.float32)
     pad_h, pad_w = GH - OH, GW - OW
-
     inits = [
         _make_int64_init('sl_st', [0, 0, 0, 0]),
         _make_int64_init('sl_en', [1, 10, IH, IW]),
-        numpy_helper.from_array(non_bg_ch_mask, 'nbg_ch'),
+        numpy_helper.from_array(non_bg_ch, 'nbg_ch'),
         _make_int64_init('rs_ax', [1]),
-        # Block-upscale via Reshape+Tile (no Resize needed)
-        _make_int64_init('shape_6d', [1, 1, IH, 1, IW, 1]),
-        _make_int64_init('tile_6d', [1, 1, 1, IH, 1, IW]),
-        _make_int64_init('shape_4d', [1, 1, OH, OW]),
+        numpy_helper.from_array(up_k, 'up_k'),
         _make_int64_init('tile_rp', [1, 1, IH, IW]),
-        numpy_helper.from_array(np.array([0, 0, 0, 0, 0, 0, pad_h, pad_w], dtype=np.int64), 'pads'),
+        numpy_helper.from_array(np.array([0,0,0,0,0,0,pad_h,pad_w], dtype=np.int64), 'pads'),
         numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'pad_cv'),
     ]
-
     nodes = [
         helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
         helper.make_node('Mul', ['cropped', 'nbg_ch'], ['non_bg_chs']),
         helper.make_node('ReduceSum', ['non_bg_chs', 'rs_ax'], ['mask_small'], keepdims=1),
-        # Block-upscale: [1,1,IH,IW] → [1,1,IH,1,IW,1] → tile → [1,1,IH,IH,IW,IW] → reshape [1,1,OH,OW]
-        helper.make_node('Reshape', ['mask_small', 'shape_6d'], ['mask_6d']),
-        helper.make_node('Tile', ['mask_6d', 'tile_6d'], ['mask_tiled']),
-        helper.make_node('Reshape', ['mask_tiled', 'shape_4d'], ['mask_big']),
-        # Tile input
+        helper.make_node('ConvTranspose', ['mask_small', 'up_k'], ['mask_big'],
+                        kernel_shape=[IH, IW], strides=[IH, IW]),
         helper.make_node('Tile', ['cropped', 'tile_rp'], ['tiled']),
         helper.make_node('Mul', ['tiled', 'mask_big'], ['result']),
         helper.make_node('Pad', ['result', 'pads', 'pad_cv'], ['output'], mode='constant'),
     ]
-
     return mk(nodes, inits)
 
 
-# =============================================================================
-# SELF-TILE SOLVER (majority color mask)
-# =============================================================================
-
 def s_self_tile_majority(td):
-    """Self-tiling with majority-color mask.
-
-    Each pixel matching the majority color becomes a full copy of input.
-    Other pixels become blocks of zeros.
-    Output size = IH*IH x IW*IW.
-    """
+    """Self-tiling with majority-color mask. Output = IH^2 x IW^2."""
     exs = get_exs(td)
     sp = fixed_shapes(td)
     if sp is None:
         return None
     (IH, IW), (OH, OW) = sp
-
     if OH != IH * IH or OW != IW * IW:
         return None
     if OH > 30 or OW > 30:
         return None
-
     for inp, out in exs:
         counts = np.bincount(inp.flatten(), minlength=10)
         majority = int(np.argmax(counts))
@@ -119,55 +82,41 @@ def s_self_tile_majority(td):
         if not np.array_equal(expected, out):
             return None
 
+    up_k = np.ones((1, 1, IH, IW), dtype=np.float32)
     pad_h, pad_w = GH - OH, GW - OW
-
     inits = [
         _make_int64_init('sl_st', [0, 0, 0, 0]),
         _make_int64_init('sl_en', [1, 10, IH, IW]),
         _make_int64_init('rs_ax_spatial', [2, 3]),
         _make_int64_init('rs_ax_channel', [1]),
         numpy_helper.from_array(np.arange(10, dtype=np.int64).reshape(1, 10, 1, 1), 'classes'),
-        # Block-upscale shapes
-        _make_int64_init('shape_6d', [1, 1, IH, 1, IW, 1]),
-        _make_int64_init('tile_6d', [1, 1, 1, IH, 1, IW]),
-        _make_int64_init('shape_4d', [1, 1, OH, OW]),
+        numpy_helper.from_array(up_k, 'up_k'),
         _make_int64_init('tile_rp', [1, 1, IH, IW]),
-        numpy_helper.from_array(np.array([0, 0, 0, 0, 0, 0, pad_h, pad_w], dtype=np.int64), 'pads'),
+        numpy_helper.from_array(np.array([0,0,0,0,0,0,pad_h,pad_w], dtype=np.int64), 'pads'),
         numpy_helper.from_array(np.array(0.0, dtype=np.float32), 'pad_cv'),
     ]
-
     nodes = [
         helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
-        # Find majority color
         helper.make_node('ReduceSum', ['cropped', 'rs_ax_spatial'], ['ch_counts'], keepdims=1),
         helper.make_node('ArgMax', ['ch_counts'], ['maj_idx'], axis=1, keepdims=1),
-        # One-hot majority
         helper.make_node('Equal', ['maj_idx', 'classes'], ['maj_eq']),
         helper.make_node('Cast', ['maj_eq'], ['maj_oh'], to=TensorProto.FLOAT),
-        # Extract majority channel as mask
         helper.make_node('Mul', ['cropped', 'maj_oh'], ['maj_channel']),
         helper.make_node('ReduceSum', ['maj_channel', 'rs_ax_channel'], ['mask_small'], keepdims=1),
-        # Block-upscale mask (Reshape+Tile, no Resize)
-        helper.make_node('Reshape', ['mask_small', 'shape_6d'], ['mask_6d']),
-        helper.make_node('Tile', ['mask_6d', 'tile_6d'], ['mask_tiled']),
-        helper.make_node('Reshape', ['mask_tiled', 'shape_4d'], ['mask_big']),
-        # Tile input
+        helper.make_node('ConvTranspose', ['mask_small', 'up_k'], ['mask_big'],
+                        kernel_shape=[IH, IW], strides=[IH, IW]),
         helper.make_node('Tile', ['cropped', 'tile_rp'], ['tiled']),
-        # Multiply
         helper.make_node('Mul', ['tiled', 'mask_big'], ['result']),
-        # Pad
         helper.make_node('Pad', ['result', 'pads', 'pad_cv'], ['output'], mode='constant'),
     ]
-
     return mk(nodes, inits)
 
 
 # =============================================================================
-# SPLIT-AND-MERGE SOLVERS
+# SPLIT-AND-MERGE
 # =============================================================================
 
 def _build_split_or_model(IH, IW, OH, OW, direction, sep_pos):
-    """Slice two halves + Add + channel-weighted ArgMax."""
     ch_weights = np.arange(10, dtype=np.float32).reshape(1, 10, 1, 1)
     pad_h, pad_w = GH - OH, GW - OW
     classes = np.arange(10, dtype=np.int64).reshape(1, 10, 1, 1)
@@ -195,7 +144,6 @@ def _build_split_or_model(IH, IW, OH, OW, direction, sep_pos):
 
 
 def _build_split_and_model(IH, IW, OH, OW, direction, sep_pos, out_color):
-    """AND of non-bg masks from two halves."""
     pad_h, pad_w = GH - OH, GW - OW
     non_bg_mask = np.ones((1,10,1,1), dtype=np.float32); non_bg_mask[0,0,0,0] = 0.0
     out_oh = np.zeros((1,10,1,1), dtype=np.float32); out_oh[0,out_color,0,0] = 1.0
@@ -228,7 +176,6 @@ def _build_split_and_model(IH, IW, OH, OW, direction, sep_pos, out_color):
 
 
 def _build_split_xor_model(IH, IW, OH, OW, direction, sep_pos, out_color):
-    """XOR of non-bg masks from two halves."""
     pad_h, pad_w = GH - OH, GW - OW
     non_bg_mask = np.ones((1,10,1,1), dtype=np.float32); non_bg_mask[0,0,0,0] = 0.0
     out_oh = np.zeros((1,10,1,1), dtype=np.float32); out_oh[0,out_color,0,0] = 1.0
@@ -323,12 +270,11 @@ def s_split_and_merge(td):
             model = _build_split_xor_model(IH, IW, OH, OW, 'vertical', sep_col, oc)
             if _validate_model(model, exs, OH, OW):
                 return model
-
     return None
 
 
 def _validate_model(model, exs, OH, OW):
-    """Quick validation: run model on all examples."""
+    """Quick validation on train+test examples."""
     import onnxruntime as ort
     import tempfile, os, onnx
     try: