Add Wave 2 solvers: 17 new analytical solvers (overlay, bbox_crop, row/col_mode_fill, fill_bg, pad_align, multi_stamp, diagonal_flip, invert_colors, majority_fill, border_extract, interior_fill, repeat_row/col, swap_colors, max_pool, crop_paste)

own-solver/neurogolf_solver/solvers/wave2.py

@@ -0,0 +1,863 @@
+#!/usr/bin/env python3
+"""Wave 2 analytical solvers — new patterns not covered by existing solvers.
+
+New solvers:
+- s_overlay_constant: input + fixed constant overlay
+- s_bbox_crop: crop to bounding box of non-bg pixels
+- s_row_mode_fill: each row filled with its mode color
+- s_col_mode_fill: each column filled with its mode color
+- s_fill_bg_with_mode: bg pixels → dominant non-bg color
+- s_pad_align: input pasted into larger canvas at fixed offset
+- s_multi_stamp: union of shifted copies
+- s_diagonal_flip: anti-diagonal transpose
+- s_invert_colors: color inversion (max_val - c)
+- s_majority_color_fill: solid fill with least common input color
+- s_border_extract: keep only border pixels
+- s_interior_fill: keep only interior pixels
+- s_repeat_row: repeat one row vertically
+- s_repeat_col: repeat one column horizontally
+- s_swap_two_colors: swap exactly two colors
+- s_max_pool_downsample: max-pool style downsampling
+"""
+
+import numpy as np
+from onnx import helper, numpy_helper, TensorProto
+from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, add_onehot_block
+from ..data_loader import get_exs, fixed_shapes
+from ..gather_helpers import _build_gather_model, _build_gather_model_with_const
+from ..constants import GH, GW
+
+
+# =============================================================================
+# SOLVER: overlay_constant
+# =============================================================================
+
+def s_overlay_constant(td):
+    """Output = input with specific fixed positions overwritten by constant colors."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    diff_positions = None
+    const_values = {}
+    
+    for inp, out in exs:
+        current_diffs = set()
+        for r in range(IH):
+            for c in range(IW):
+                if inp[r, c] != out[r, c]:
+                    current_diffs.add((r, c))
+                    key = (r, c)
+                    if key in const_values:
+                        if const_values[key] != int(out[r, c]):
+                            return None
+                    else:
+                        const_values[key] = int(out[r, c])
+        
+        if diff_positions is None:
+            diff_positions = current_diffs
+        elif diff_positions != current_diffs:
+            return None
+    
+    if not diff_positions or len(diff_positions) == 0:
+        return None
+    if len(diff_positions) > IH * IW * 0.8:
+        return None
+    
+    idx = np.zeros((OH, OW, 2), dtype=np.int64)
+    cst = np.full((OH, OW), -1, dtype=np.int64)
+    
+    for r in range(OH):
+        for c in range(OW):
+            if (r, c) in const_values:
+                idx[r, c] = [-1, -1]
+                cst[r, c] = const_values[(r, c)]
+            else:
+                idx[r, c] = [r, c]
+    
+    return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+
+
+# =============================================================================
+# SOLVER: bbox_crop
+# =============================================================================
+
+def s_bbox_crop(td):
+    """Crop input to the bounding box of non-background pixels."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if OH >= IH and OW >= IW:
+        return None
+    
+    for bg_color in range(10):
+        ok = True
+        for inp, out in exs:
+            rows = np.any(inp != bg_color, axis=1)
+            cols = np.any(inp != bg_color, axis=0)
+            if not np.any(rows) or not np.any(cols):
+                ok = False
+                break
+            r_min, r_max = np.where(rows)[0][[0, -1]]
+            c_min, c_max = np.where(cols)[0][[0, -1]]
+            cropped = inp[r_min:r_max+1, c_min:c_max+1]
+            if cropped.shape != out.shape or not np.array_equal(cropped, out):
+                ok = False
+                break
+        if ok:
+            inp0 = exs[0][0]
+            rows = np.any(inp0 != bg_color, axis=1)
+            cols = np.any(inp0 != bg_color, axis=0)
+            r_min = int(np.where(rows)[0][0])
+            c_min = int(np.where(cols)[0][0])
+            
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    idx[r, c] = [r + r_min, c + c_min]
+            return _build_gather_model(OH, OW, idx)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: row_mode_fill
+# =============================================================================
+
+def s_row_mode_fill(td):
+    """Each row of output = solid fill of the mode color of that input row."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    for inp, out in exs:
+        for r in range(IH):
+            counts = np.bincount(inp[r, :], minlength=10)
+            mode_color = int(np.argmax(counts))
+            if not np.all(out[r, :] == mode_color):
+                return None
+    
+    pad_h, pad_w = GH - IH, GW - IW
+    inits = [
+        _make_int64_init('sl_st', [0, 0, 0, 0]),
+        _make_int64_init('sl_en', [1, 10, IH, IW]),
+        _make_int64_init('rs_ax', [3]),
+        _make_int64_init('tile_rp', [1, 1, 1, IW]),
+    ]
+    nodes = [
+        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+        helper.make_node('ReduceSum', ['cropped', 'rs_ax'], ['row_sums'], keepdims=1),
+        helper.make_node('ArgMax', ['row_sums'], ['row_modes'], axis=1, keepdims=1),
+        helper.make_node('Tile', ['row_modes', 'tile_rp'], ['tiled']),
+    ]
+    add_onehot_block(nodes, inits, 'tiled', 'oh_out')
+    nodes.append(_build_pad_node('oh_out', 'output', pad_h, pad_w, inits))
+    return mk(nodes, inits)
+
+
+# =============================================================================
+# SOLVER: col_mode_fill
+# =============================================================================
+
+def s_col_mode_fill(td):
+    """Each column of output = solid fill of the mode color of that input column."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    for inp, out in exs:
+        for c in range(IW):
+            counts = np.bincount(inp[:, c], minlength=10)
+            mode_color = int(np.argmax(counts))
+            if not np.all(out[:, c] == mode_color):
+                return None
+    
+    pad_h, pad_w = GH - IH, GW - IW
+    inits = [
+        _make_int64_init('sl_st', [0, 0, 0, 0]),
+        _make_int64_init('sl_en', [1, 10, IH, IW]),
+        _make_int64_init('rs_ax', [2]),
+        _make_int64_init('tile_rp', [1, 1, IH, 1]),
+    ]
+    nodes = [
+        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+        helper.make_node('ReduceSum', ['cropped', 'rs_ax'], ['col_sums'], keepdims=1),
+        helper.make_node('ArgMax', ['col_sums'], ['col_modes'], axis=1, keepdims=1),
+        helper.make_node('Tile', ['col_modes', 'tile_rp'], ['tiled']),
+    ]
+    add_onehot_block(nodes, inits, 'tiled', 'oh_out')
+    nodes.append(_build_pad_node('oh_out', 'output', pad_h, pad_w, inits))
+    return mk(nodes, inits)
+
+
+# =============================================================================
+# SOLVER: fill_bg_with_mode
+# =============================================================================
+
+def s_fill_bg_with_mode(td):
+    """Background pixels replaced by the most common non-background color."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    for bg_color in range(10):
+        ok = True
+        for inp, out in exs:
+            non_bg = inp[inp != bg_color]
+            if len(non_bg) == 0:
+                ok = False
+                break
+            counts = np.bincount(non_bg, minlength=10)
+            mode_color = int(np.argmax(counts))
+            expected = inp.copy()
+            expected[inp == bg_color] = mode_color
+            if not np.array_equal(expected, out):
+                ok = False
+                break
+        
+        if ok:
+            pad_h, pad_w = GH - IH, GW - IW
+            non_bg_mask = np.ones((1, 10, 1, 1), dtype=np.float32)
+            non_bg_mask[0, bg_color, 0, 0] = 0.0
+            
+            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_mask, 'non_bg_mask'),
+                _make_int64_init('rs_ax', [2, 3]),
+                _make_int64_init('tile_rp', [1, 1, IH, IW]),
+                numpy_helper.from_array(np.float32(0.5), 'half'),
+                _make_int64_init('bg_st', [0, bg_color, 0, 0]),
+                _make_int64_init('bg_en', [1, bg_color + 1, IH, IW]),
+            ]
+            nodes = [
+                helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+                helper.make_node('Mul', ['cropped', 'non_bg_mask'], ['non_bg_only']),
+                helper.make_node('ReduceSum', ['non_bg_only', 'rs_ax'], ['color_counts'], keepdims=1),
+                helper.make_node('ArgMax', ['color_counts'], ['mode_idx'], axis=1, keepdims=1),
+                helper.make_node('Tile', ['mode_idx', 'tile_rp'], ['mode_tiled']),
+            ]
+            add_onehot_block(nodes, inits, 'mode_tiled', 'mode_oh')
+            nodes.extend([
+                helper.make_node('Slice', ['cropped', 'bg_st', 'bg_en'], ['bg_ch']),
+                helper.make_node('Greater', ['bg_ch', 'half'], ['is_bg']),
+                helper.make_node('Where', ['is_bg', 'mode_oh', 'cropped'], ['filled']),
+            ])
+            nodes.append(_build_pad_node('filled', 'output', pad_h, pad_w, inits))
+            return mk(nodes, inits)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: pad_align
+# =============================================================================
+
+def s_pad_align(td):
+    """Input pasted into larger canvas at fixed offset, rest = bg_color."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if OH <= IH and OW <= IW:
+        return None
+    if OH > 30 or OW > 30:
+        return None
+    
+    for bg_color in range(10):
+        for r0 in range(OH - IH + 1):
+            for c0 in range(OW - IW + 1):
+                ok = True
+                for inp, out in exs:
+                    expected = np.full((OH, OW), bg_color, dtype=np.int64)
+                    expected[r0:r0+IH, c0:c0+IW] = inp
+                    if not np.array_equal(expected, out):
+                        ok = False
+                        break
+                if ok:
+                    idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                    cst = np.full((OH, OW), -1, dtype=np.int64)
+                    for r in range(OH):
+                        for c in range(OW):
+                            if r0 <= r < r0 + IH and c0 <= c < c0 + IW:
+                                idx[r, c] = [r - r0, c - c0]
+                            else:
+                                idx[r, c] = [-1, -1]
+                                cst[r, c] = bg_color
+                    return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+# =============================================================================
+# SOLVER: multi_stamp
+# =============================================================================
+
+def s_multi_stamp(td):
+    """Output = union of multiple shifted copies of input at fixed offsets."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    def apply_stamps(inp, offsets, bg=0):
+        result = np.full((IH, IW), bg, dtype=np.int64)
+        for dr, dc in offsets:
+            for r in range(IH):
+                for c in range(IW):
+                    sr, sc = r - dr, c - dc
+                    if 0 <= sr < IH and 0 <= sc < IW and inp[sr, sc] != bg:
+                        result[r, c] = inp[sr, sc]
+        return result
+    
+    for bg in range(2):
+        small_range = range(-min(IH, 8), min(IH, 8) + 1)
+        for dr1 in small_range:
+            for dc1 in small_range:
+                if dr1 == 0 and dc1 == 0:
+                    continue
+                offsets = [(0, 0), (dr1, dc1)]
+                if all(np.array_equal(apply_stamps(inp, offsets, bg), out) 
+                       for inp, out in exs):
+                    return _build_stamp_model(IH, IW, offsets, bg)
+    return None
+
+
+def _build_stamp_model(IH, IW, offsets, bg):
+    """Build ONNX model for multi-stamp."""
+    OH, OW = IH, IW
+    pad_h, pad_w = GH - IH, GW - IW
+    n_stamps = len(offsets)
+    
+    inits = [
+        _make_int64_init('sl_st', [0, 0, 0, 0]),
+        _make_int64_init('sl_en', [1, 10, IH, IW]),
+    ]
+    nodes = [
+        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+    ]
+    
+    stamp_names = []
+    for i, (dr, dc) in enumerate(offsets):
+        flat_idx = np.zeros((IH * IW,), dtype=np.int64)
+        mask = np.zeros((1, 1, IH, IW), dtype=np.float32)
+        
+        for r in range(IH):
+            for c in range(IW):
+                sr, sc = r - dr, c - dc
+                if 0 <= sr < IH and 0 <= sc < IW:
+                    flat_idx[r * IW + c] = sr * IW + sc
+                    mask[0, 0, r, c] = 1.0
+        
+        idx_name = f'sidx_{i}'
+        mask_name = f'smask_{i}'
+        inits.append(numpy_helper.from_array(flat_idx, idx_name))
+        inits.append(numpy_helper.from_array(mask, mask_name))
+        inits.append(_make_int64_init(f'sfs_{i}', [1, 10, IH * IW]))
+        inits.append(_make_int64_init(f'sos_{i}', [1, 10, IH, IW]))
+        
+        nodes.extend([
+            helper.make_node('Reshape', ['cropped', f'sfs_{i}'], [f'sflat_{i}']),
+            helper.make_node('Gather', [f'sflat_{i}', idx_name], [f'sg_{i}'], axis=2),
+            helper.make_node('Reshape', [f'sg_{i}', f'sos_{i}'], [f'sraw_{i}']),
+            helper.make_node('Mul', [f'sraw_{i}', mask_name], [f'stamp_{i}']),
+        ])
+        stamp_names.append(f'stamp_{i}')
+    
+    # Sum all stamps
+    if len(stamp_names) == 2:
+        nodes.append(helper.make_node('Add', stamp_names, ['sum_stamps']))
+    else:
+        nodes.append(helper.make_node('Add', [stamp_names[0], stamp_names[1]], ['sum_01']))
+        cur = 'sum_01'
+        for i in range(2, len(stamp_names)):
+            nxt = f'sum_0{i}' if i < len(stamp_names) - 1 else 'sum_stamps'
+            nodes.append(helper.make_node('Add', [cur, stamp_names[i]], [nxt]))
+            cur = nxt
+    
+    # Weight channels so non-bg colors always win over bg
+    ch_weights = np.arange(10, dtype=np.float32).reshape(1, 10, 1, 1)
+    inits.append(numpy_helper.from_array(ch_weights, 'ch_w'))
+    nodes.append(helper.make_node('Mul', ['sum_stamps', 'ch_w'], ['weighted']))
+    nodes.append(helper.make_node('ArgMax', ['weighted'], ['am'], axis=1, keepdims=1))
+    add_onehot_block(nodes, inits, 'am', 'oh_out')
+    nodes.append(_build_pad_node('oh_out', 'output', pad_h, pad_w, inits))
+    return mk(nodes, inits)
+
+
+# =============================================================================
+# SOLVER: diagonal_flip
+# =============================================================================
+
+def s_diagonal_flip(td):
+    """Flip along anti-diagonal: out[r,c] = inp[IW-1-c, IH-1-r]."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if OH != IW or OW != IH:
+        return None
+    
+    ok = True
+    for inp, out in exs:
+        expected = np.zeros((OH, OW), dtype=np.int64)
+        for r in range(OH):
+            for c in range(OW):
+                expected[r, c] = inp[IW - 1 - c, IH - 1 - r]
+        if not np.array_equal(expected, out):
+            ok = False
+            break
+    
+    if ok:
+        idx = np.zeros((OH, OW, 2), dtype=np.int64)
+        for r in range(OH):
+            for c in range(OW):
+                idx[r, c] = [IW - 1 - c, IH - 1 - r]
+        return _build_gather_model(OH, OW, idx)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: invert_colors
+# =============================================================================
+
+def s_invert_colors(td):
+    """Color inversion: out = max_val - inp."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    for max_val in range(1, 10):
+        ok = True
+        for inp, out in exs:
+            expected = max_val - inp
+            if np.any(expected < 0) or np.any(expected >= 10):
+                ok = False
+                break
+            if not np.array_equal(out, expected):
+                ok = False
+                break
+        if ok:
+            W = np.zeros((10, 10, 1, 1), dtype=np.float32)
+            for c in range(10):
+                dst = max_val - c
+                if 0 <= dst < 10:
+                    W[dst, c, 0, 0] = 1.0
+            inits = [numpy_helper.from_array(W, 'W')]
+            nodes = [helper.make_node('Conv', ['input', 'W'], ['output'], kernel_shape=[1, 1])]
+            return mk(nodes, inits)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: majority_color_fill (solid fill with least common color)
+# =============================================================================
+
+def s_majority_color_fill(td):
+    """Output = solid grid filled with the least common present color of input."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    
+    # Check if output is always a solid single color
+    for _, out in exs:
+        if not np.all(out == out[0, 0]):
+            return None
+    
+    # Check if color is the LEAST common present color
+    ok = True
+    for inp, out in exs:
+        out_color = int(out[0, 0])
+        counts = np.bincount(inp.flatten(), minlength=10)
+        present = np.where(counts > 0)[0]
+        if len(present) < 2:
+            ok = False
+            break
+        min_count_color = present[np.argmin(counts[present])]
+        if min_count_color != out_color:
+            ok = False
+            break
+    
+    if ok and len(exs) >= 2:
+        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', [2, 3]),
+            _make_int64_init('tile_rp', [1, 1, OH, OW]),
+            numpy_helper.from_array(np.float32(-1.0), 'neg_one'),
+            numpy_helper.from_array(np.float32(0.5), 'half'),
+        ]
+        nodes = [
+            helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+            helper.make_node('ReduceSum', ['cropped', 'rs_ax'], ['counts'], keepdims=1),
+            # Negate so argmax gives argmin
+            helper.make_node('Mul', ['counts', 'neg_one'], ['neg_counts']),
+            # Mask out zero-count channels (so they don't win)
+            helper.make_node('Greater', ['counts', 'half'], ['present']),
+            helper.make_node('Cast', ['present'], ['present_f'], to=TensorProto.FLOAT),
+            helper.make_node('Mul', ['neg_counts', 'present_f'], ['masked_neg']),
+            helper.make_node('ArgMax', ['masked_neg'], ['min_idx'], axis=1, keepdims=1),
+            helper.make_node('Tile', ['min_idx', 'tile_rp'], ['tiled']),
+        ]
+        add_onehot_block(nodes, inits, 'tiled', 'oh_out')
+        nodes.append(_build_pad_node('oh_out', 'output', pad_h, pad_w, inits))
+        return mk(nodes, inits)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: border_extract
+# =============================================================================
+
+def s_border_extract(td):
+    """Output = only the border pixels of input, interior → bg."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    if IH < 3 or IW < 3:
+        return None
+    
+    for bg in range(10):
+        ok = True
+        for inp, out in exs:
+            expected = np.full((IH, IW), bg, dtype=np.int64)
+            expected[0, :] = inp[0, :]
+            expected[-1, :] = inp[-1, :]
+            expected[:, 0] = inp[:, 0]
+            expected[:, -1] = inp[:, -1]
+            if not np.array_equal(expected, out):
+                ok = False
+                break
+        if ok:
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            cst = np.full((OH, OW), -1, dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    if r == 0 or r == IH - 1 or c == 0 or c == IW - 1:
+                        idx[r, c] = [r, c]
+                    else:
+                        idx[r, c] = [-1, -1]
+                        cst[r, c] = bg
+            return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+# =============================================================================
+# SOLVER: interior_fill
+# =============================================================================
+
+def s_interior_fill(td):
+    """Keep only interior pixels, border → bg. Same size output."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    if IH < 3 or IW < 3:
+        return None
+    
+    for b in range(1, min(IH, IW) // 2):
+        for bg in range(10):
+            ok = True
+            for inp, out in exs:
+                expected = np.full((IH, IW), bg, dtype=np.int64)
+                expected[b:IH-b, b:IW-b] = inp[b:IH-b, b:IW-b]
+                if not np.array_equal(expected, out):
+                    ok = False
+                    break
+            if ok:
+                idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                cst = np.full((OH, OW), -1, dtype=np.int64)
+                for r in range(OH):
+                    for c in range(OW):
+                        if b <= r < IH - b and b <= c < IW - b:
+                            idx[r, c] = [r, c]
+                        else:
+                            idx[r, c] = [-1, -1]
+                            cst[r, c] = bg
+                return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+# =============================================================================
+# SOLVER: repeat_row
+# =============================================================================
+
+def s_repeat_row(td):
+    """Output = one specific row of input repeated vertically to fill grid."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if IW != OW:
+        return None
+    
+    for src_row in range(IH):
+        ok = True
+        for inp, out in exs:
+            expected = np.tile(inp[src_row:src_row+1, :], (OH, 1))
+            if not np.array_equal(expected, out):
+                ok = False
+                break
+        if ok:
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    idx[r, c] = [src_row, c]
+            return _build_gather_model(OH, OW, idx)
+    return None
+
+
+# =============================================================================
+# SOLVER: repeat_col
+# =============================================================================
+
+def s_repeat_col(td):
+    """Output = one specific column of input repeated horizontally to fill grid."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if IH != OH:
+        return None
+    
+    for src_col in range(IW):
+        ok = True
+        for inp, out in exs:
+            expected = np.tile(inp[:, src_col:src_col+1], (1, OW))
+            if not np.array_equal(expected, out):
+                ok = False
+                break
+        if ok:
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    idx[r, c] = [r, src_col]
+            return _build_gather_model(OH, OW, idx)
+    return None
+
+
+# =============================================================================
+# SOLVER: swap_two_colors
+# =============================================================================
+
+def s_swap_two_colors(td):
+    """Swap exactly two colors A and B, all other colors unchanged."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    
+    swaps = {}
+    for inp, out in exs:
+        for iv, ov in zip(inp.flat, out.flat):
+            iv, ov = int(iv), int(ov)
+            if iv != ov:
+                if iv in swaps:
+                    if swaps[iv] != ov:
+                        return None
+                else:
+                    swaps[iv] = ov
+    
+    if len(swaps) != 2:
+        return None
+    
+    items = list(swaps.items())
+    if items[0][1] != items[1][0] or items[1][1] != items[0][0]:
+        return None
+    
+    a, b = items[0][0], items[0][1]
+    for inp, out in exs:
+        expected = inp.copy()
+        expected[inp == a] = b
+        expected[inp == b] = a
+        if not np.array_equal(expected, out):
+            return None
+    
+    gather_ch = np.arange(10, dtype=np.int32)
+    gather_ch[a] = b
+    gather_ch[b] = a
+    inits = [numpy_helper.from_array(gather_ch, 'gi')]
+    nodes = [helper.make_node('Gather', ['input', 'gi'], ['output'], axis=1)]
+    return mk(nodes, inits)
+
+
+# =============================================================================
+# SOLVER: max_pool_downsample
+# =============================================================================
+
+def s_max_pool_downsample(td):
+    """Non-overlapping max-pool: output = max color value in each block."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if OH >= IH or OW >= IW:
+        return None
+    if IH % OH != 0 or IW % OW != 0:
+        return None
+    
+    bh, bw = IH // OH, IW // OW
+    if bh < 2 and bw < 2:
+        return None
+    
+    # Check if simple max works
+    ok = True
+    for inp, out in exs:
+        for r in range(OH):
+            for c in range(OW):
+                block = inp[r*bh:(r+1)*bh, c*bw:(c+1)*bw]
+                if out[r, c] != int(np.max(block)):
+                    ok = False
+                    break
+            if not ok:
+                break
+        if not ok:
+            break
+    
+    if ok:
+        pad_h, pad_w = GH - OH, GW - OW
+        ch_weights = np.arange(10, dtype=np.float32).reshape(1, 10, 1, 1)
+        inits = [
+            _make_int64_init('sl_st', [0, 0, 0, 0]),
+            _make_int64_init('sl_en', [1, 10, IH, IW]),
+            numpy_helper.from_array(ch_weights, 'ch_w'),
+        ]
+        nodes = [
+            helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
+            helper.make_node('MaxPool', ['cropped'], ['pooled'],
+                           kernel_shape=[bh, bw], strides=[bh, bw]),
+            helper.make_node('Mul', ['pooled', 'ch_w'], ['weighted']),
+            helper.make_node('ArgMax', ['weighted'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(_build_pad_node('oh_out', 'output', pad_h, pad_w, inits))
+        return mk(nodes, inits)
+    
+    return None
+
+
+# =============================================================================
+# SOLVER: crop_paste (limited search for speed)
+# =============================================================================
+
+def s_crop_paste(td):
+    """Crop from one fixed position, paste at another. Limited search."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if (IH, IW) != (OH, OW):
+        return None
+    if IH < 3 or IW < 3:
+        return None
+    
+    for bg in range(2):
+        out0 = exs[0][1]
+        rows = np.any(out0 != bg, axis=1)
+        cols = np.any(out0 != bg, axis=0)
+        if not np.any(rows) or not np.any(cols):
+            continue
+        dr = int(np.where(rows)[0][0])
+        dc = int(np.where(cols)[0][0])
+        dr_end = int(np.where(rows)[0][-1]) + 1
+        dc_end = int(np.where(cols)[0][-1]) + 1
+        ch = dr_end - dr
+        cw = dc_end - dc
+        
+        if ch < 2 or cw < 2 or ch >= IH or cw >= IW:
+            continue
+        
+        for sr in range(IH - ch + 1):
+            for sc in range(IW - cw + 1):
+                if sr == dr and sc == dc:
+                    continue
+                ok = True
+                for inp, out in exs:
+                    expected = np.full((OH, OW), bg, dtype=np.int64)
+                    expected[dr:dr+ch, dc:dc+cw] = inp[sr:sr+ch, sc:sc+cw]
+                    if not np.array_equal(expected, out):
+                        ok = False
+                        break
+                if ok:
+                    idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                    cst = np.full((OH, OW), -1, dtype=np.int64)
+                    for r in range(OH):
+                        for c in range(OW):
+                            if dr <= r < dr + ch and dc <= c < dc + cw:
+                                idx[r, c] = [sr + (r - dr), sc + (c - dc)]
+                            else:
+                                idx[r, c] = [-1, -1]
+                                cst[r, c] = bg
+                    return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+# =============================================================================
+# Collect all Wave 2 solvers for registration
+# =============================================================================
+
+WAVE2_SOLVERS = [
+    ('overlay_constant', s_overlay_constant),
+    ('bbox_crop', s_bbox_crop),
+    ('row_mode_fill', s_row_mode_fill),
+    ('col_mode_fill', s_col_mode_fill),
+    ('fill_bg_with_mode', s_fill_bg_with_mode),
+    ('pad_align', s_pad_align),
+    ('multi_stamp', s_multi_stamp),
+    ('diagonal_flip', s_diagonal_flip),
+    ('invert_colors', s_invert_colors),
+    ('majority_color_fill', s_majority_color_fill),
+    ('border_extract', s_border_extract),
+    ('interior_fill', s_interior_fill),
+    ('repeat_row', s_repeat_row),
+    ('repeat_col', s_repeat_col),
+    ('swap_two_colors', s_swap_two_colors),
+    ('max_pool_downsample', s_max_pool_downsample),
+    ('crop_paste', s_crop_paste),
+]