medal-solvers/build_task100_onnx.py

"""Build optimized ONNX for Task 100.

Rule (266/266 verified):
- Input has 2 colored rectangles on bg=0
- Output is 2x2 filled with the color of the LARGER rectangle (by bounding box area)

ONNX approach:
- For each color 1-9: compute bounding box area
- ArgMax to find color with largest area
- Output: 2x2 of that color + bg for rest
"""
import sys
sys.path.insert(0, '/app/repo/medal-solvers')

from onnx import TensorProto
import numpy as np
from onnx_builder import OnnxBuilder


def build_task100():
    b = OnnxBuilder()
    const, nd = b.const, b.nd

    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_zero', [0.0])
    const('c_big', [1000.0])
    const('c_neg_big', [-1000.0])
    const('c_1_5', [1.5])
    
    const('axes2', [2], 'i')
    const('axes3', [3], 'i')
    const('axes23', [2, 3], 'i')
    const('axes1', [1], 'i')
    
    const('shape_1_9', [1, 9], 'i')
    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    const('shape_1', [1], 'i')
    
    # Row/col grids [1,1,30,30]
    row_grid = np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1) * np.ones((1, 1, 1, 30), dtype=np.float32)
    col_grid = np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30) * np.ones((1, 1, 30, 1), dtype=np.float32)
    const('row_grid', row_grid)
    const('col_grid', col_grid)
    
    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])
    const('idx_0', [0], 'i')
    
    # For each color 1-9, compute bounding box area
    areas = []  # list of [1,1,1,1] tensors
    
    for c in range(1, 10):
        const(f'idx_{c}', [c], 'i')
        ch = nd('Gather', ['input', f'idx_{c}'], [[1, 1, 30, 30]], axis=1)
        
        # Has any pixel of this color?
        ch_sum = nd('ReduceSum', [ch, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
        has_color = nd('Greater', [ch_sum, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
        has_color_f = nd('Cast', [has_color], [[1, 1, 1, 1]], to=1)
        
        # Bounding box: min/max row and col where ch > 0
        not_ch = nd('Sub', ['c_one', ch], [[1, 1, 30, 30]])
        
        # r_min
        rg_ch = nd('Add', [nd('Mul', [ch, 'row_grid'], [[1, 1, 30, 30]]),
                           nd('Mul', [not_ch, 'c_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]])
        r_min = nd('ReduceMin', [rg_ch, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
        
        # r_max
        rg_ch_max = nd('Add', [nd('Mul', [ch, 'row_grid'], [[1, 1, 30, 30]]),
                               nd('Mul', [not_ch, 'c_neg_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]])
        r_max = nd('ReduceMax', [rg_ch_max, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
        
        # c_min
        cg_ch = nd('Add', [nd('Mul', [ch, 'col_grid'], [[1, 1, 30, 30]]),
                           nd('Mul', [not_ch, 'c_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]])
        c_min = nd('ReduceMin', [cg_ch, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
        
        # c_max
        cg_ch_max = nd('Add', [nd('Mul', [ch, 'col_grid'], [[1, 1, 30, 30]]),
                               nd('Mul', [not_ch, 'c_neg_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]])
        c_max = nd('ReduceMax', [cg_ch_max, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
        
        # Area = (r_max - r_min + 1) * (c_max - c_min + 1), gated by has_color
        height = nd('Add', [nd('Sub', [r_max, r_min], [[1, 1, 1, 1]]), 'c_one'], [[1, 1, 1, 1]])
        width = nd('Add', [nd('Sub', [c_max, c_min], [[1, 1, 1, 1]]), 'c_one'], [[1, 1, 1, 1]])
        area = nd('Mul', [height, width], [[1, 1, 1, 1]])
        area_gated = nd('Mul', [area, has_color_f], [[1, 1, 1, 1]])
        areas.append(area_gated)
    
    # Concatenate areas to [1, 9, 1, 1] and find argmax
    areas_cat = nd('Concat', areas, [[1, 9, 1, 1]], axis=1)
    # ArgMax over axis 1 → index 0-8 (corresponds to color 1-9)
    best_idx = nd('ArgMax', [areas_cat], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    best_idx_1d = nd('Reshape', [best_idx, 'shape_1'], [([1], TensorProto.INT64)])
    
    # Add 1 to get actual color index (since areas are for colors 1-9, idx 0 = color 1)
    const('one_i', [1], 'i')
    color_idx = nd('Add', [best_idx_1d, 'one_i'], [([1], TensorProto.INT64)])
    
    # Build output: 2x2 of selected color at (0,0)-(1,1), bg elsewhere within grid
    # First get the active grid area from input
    active_sum = nd('ReduceSum', ['input', 'axes1'], [[1, 1, 30, 30]], keepdims=1)
    active_b = nd('Greater', [active_sum, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
    active = nd('Cast', [active_b], [[1, 1, 30, 30]], to=1)
    
    # 2x2 mask at top-left
    row_lt_2 = nd('Less', ['row_grid', 'c_1_5'], [([1, 1, 30, 30], TensorProto.BOOL)])
    col_lt_2 = nd('Less', ['col_grid', 'c_1_5'], [([1, 1, 30, 30], TensorProto.BOOL)])
    box_2x2 = nd('And', [row_lt_2, col_lt_2], [([1, 1, 30, 30], TensorProto.BOOL)])
    box_2x2_f = nd('Cast', [box_2x2], [[1, 1, 30, 30]], to=1)
    
    # Color the 2x2 with selected color
    color_oh = nd('OneHot', [color_idx, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    color_oh_4d = nd('Reshape', [color_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    colored_box = nd('Mul', [box_2x2_f, color_oh_4d], [[1, 10, 30, 30]])
    
    # No need for bg since output grid is exactly 2x2 and all cells are colored
    # Actually we DO need the output to match expected format:
    # Output grid is 2x2. In 30x30 one-hot: only cells (0,0),(0,1),(1,0),(1,1) should be non-zero
    # Those 4 cells should have the selected color channel = 1
    # All other cells: all channels = 0 (outside grid)
    
    return b.finish('task100', last_tensor=colored_box)


if __name__ == '__main__':
    import os, json, math, types
    import numpy as np
    import onnx
    import onnxruntime as ort
    
    print("Building Task 100 ONNX model...")
    model = build_task100()
    
    os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True)
    output_path = '/app/repo/medal-solvers/optimized/task100.onnx'
    onnx.save(model, output_path)
    
    fsize = os.path.getsize(output_path)
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  File size: {fsize:,} bytes (limit: 1,509,949)")
    
    # Quick validation
    sess = ort.InferenceSession(output_path)
    with open('/app/task-data/task100.json') as f:
        data = json.load(f)
    
    all_examples = data['train'] + data['test'] + data['arc-gen']
    right_count, wrong_count = 0, 0
    for i, ex in enumerate(all_examples):
        inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(ex['input']):
            for ci, v in enumerate(row):
                if r < 30 and ci < 30:
                    inp[0][v][r][ci] = 1.0
        result = sess.run(['output'], {'input': inp})
        out = (result[0] > 0.0).astype(float)
        exp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(ex['output']):
            for ci, v in enumerate(row):
                if r < 30 and ci < 30:
                    exp[0][v][r][ci] = 1.0
        if np.array_equal(out, exp):
            right_count += 1
        else:
            wrong_count += 1
            if wrong_count <= 2:
                diff = np.where(out != exp)
                print(f"  FAIL {i}: {len(diff[0])} diffs")
    
    print(f"\nQuick validation: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")
    
    if wrong_count == 0:
        # Official verification
        print("\nRunning official verification...")
        sys.path.insert(0, '/app/repo/medal-solvers')
        mock_ipython = types.ModuleType('IPython')
        mock_display = types.ModuleType('IPython.display')
        mock_display.display = lambda *a, **k: None
        mock_display.FileLink = lambda x: x
        mock_ipython.display = mock_display
        sys.modules['IPython'] = mock_ipython
        sys.modules['IPython.display'] = mock_display
        import matplotlib; matplotlib.use('Agg')
        import neurogolf_utils
        neurogolf_utils._NEUROGOLF_DIR = '/app/task-data/'
        
        sanitized = neurogolf_utils.sanitize_model(onnx.load(output_path))
        options = ort.SessionOptions()
        options.enable_profiling = True
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
        options.profile_file_prefix = '100'
        sess_off = ort.InferenceSession(sanitized.SerializeToString(), options)
        
        examples = neurogolf_utils.load_examples(100)
        r1, w1, _ = neurogolf_utils.verify_subset(sess_off, examples['train'] + examples['test'])
        r2, w2, _ = neurogolf_utils.verify_subset(sess_off, examples['arc-gen'])
        print(f"  ARC-AGI: {r1} pass, {w1} fail")
        print(f"  ARC-GEN: {r2} pass, {w2} fail")
        
        if w1 == 0 and w2 == 0:
            memory, params = neurogolf_utils.score_network(sanitized, sess_off.end_profiling())
            if memory and params:
                pts = max(1.0, 25.0 - math.log(max(1.0, memory + params)))
                print(f"  Score: {pts:.3f} pts (memory={memory}, params={params})")