medal-solvers/build_task062_onnx.py

"""Build OPTIMIZED ONNX model for Task 062 - v3 minimal memory.

Key optimization: avoid [1,1,10,10] intermediates for edge computation.
Use ReduceSum projections (row→[1,1,10,1], col→[1,1,1,10]) then ArgMax/ArgMin
on 1D vectors instead of masked 2D grids.
"""
import sys, os
import numpy as np
import onnx
from onnx import helper, numpy_helper, TensorProto
import math


def build_task062_v3():
    nodes, inits, vis = [], [], []
    counter = [0]

    def nm():
        counter[0] += 1
        return f"t{counter[0]}"

    def const(name, val, dtype='f'):
        arr = np.array(val, dtype=np.float32 if dtype == 'f' else np.int64)
        inits.append(numpy_helper.from_array(arr, name))

    def vi(name, shape, dt=TensorProto.FLOAT):
        vis.append(helper.make_tensor_value_info(name, dt, shape))

    def nd(op, ins, outs_shapes, **kwargs):
        out_names = []
        for sd in outs_shapes:
            if isinstance(sd, tuple):
                shape, dt = sd
            else:
                shape, dt = sd, TensorProto.FLOAT
            n = nm()
            vi(n, shape, dt)
            out_names.append(n)
        nodes.append(helper.make_node(op, ins, out_names, **kwargs))
        return out_names[0] if len(out_names) == 1 else out_names

    # === CONSTANTS ===
    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_zero', [0.0])
    const('c_two', [2.0])
    const('c_big', [100.0])
    const('c_neg_big', [-100.0])

    const('sl_start', [0, 0, 0, 0], 'i')
    const('sl_end', [1, 10, 10, 10], 'i')
    const('sl_axes', [0, 1, 2, 3], 'i')
    const('axes1', [1], 'i')
    const('axes2', [2], 'i')
    const('axes3', [3], 'i')
    const('axes23', [2, 3], 'i')

    # Grids
    const('row_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1))
    const('col_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10))
    # 1D grids for ArgMax/Min on projections
    const('row_1d', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1))
    const('col_1d', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10))

    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    const('shape_1', [1], 'i')

    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])

    const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i')
    const('pad_val_zero', [0.0])

    const('nobg_no2_mask', np.array([[[[0, 1, 0, 1, 1, 1, 1, 1, 1, 1]]]], dtype=np.float32).reshape(1, 10, 1, 1))
    const('ch2_start', [0, 2, 0, 0], 'i')
    const('ch2_end', [1, 3, 10, 10], 'i')
    const('idx_3', [3], 'i')
    const('ones_1_1_1_10', np.ones((1, 1, 1, 10), dtype=np.float32))
    const('ones_1_1_10_1', np.ones((1, 1, 10, 1), dtype=np.float32))
    const('c_0f', [0.0])
    const('c_9_0', [9.0])

    # === STEP 1: Slice to 10x10 ===
    inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]])

    # === STEP 2: Find main color ===
    ch_sums = nd('ReduceSum', [inp10, 'axes23'], [[1, 10, 1, 1]], keepdims=1)
    ch_sums_m = nd('Mul', [ch_sums, 'nobg_no2_mask'], [[1, 10, 1, 1]])
    mc_idx = nd('ArgMax', [ch_sums_m], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1)
    mc_1d = nd('Reshape', [mc_idx, 'shape_1'], [([1], TensorProto.INT64)])

    # === STEP 3: Main mask & axis mask ===
    mc_oh = nd('OneHot', [mc_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    mc_oh_4d = nd('Reshape', [mc_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    main_sel = nd('Mul', [inp10, mc_oh_4d], [[1, 10, 10, 10]])
    main_mask = nd('ReduceSum', [main_sel, 'axes1'], [[1, 1, 10, 10]], keepdims=1)
    axis_mask = nd('Slice', [inp10, 'ch2_start', 'ch2_end', 'sl_axes'], [[1, 1, 10, 10]])

    # === STEP 4: Row/col projections for centroids and edges ===
    # Main row projection [1,1,10,1]: sum over cols
    main_rp = nd('ReduceSum', [main_mask, 'axes3'], [[1, 1, 10, 1]], keepdims=1)
    # Main col projection [1,1,1,10]: sum over rows
    main_cp = nd('ReduceSum', [main_mask, 'axes2'], [[1, 1, 1, 10]], keepdims=1)
    # Axis row projection [1,1,10,1]
    axis_rp = nd('ReduceSum', [axis_mask, 'axes3'], [[1, 1, 10, 1]], keepdims=1)
    # Axis col projection [1,1,1,10]
    axis_cp = nd('ReduceSum', [axis_mask, 'axes2'], [[1, 1, 1, 10]], keepdims=1)

    # === STEP 5: Centroids from projections ===
    main_total = nd('ReduceSum', [main_rp, 'axes2'], [[1, 1, 1, 1]], keepdims=1)
    # Main centroid row = sum(main_rp * row_1d) / main_total
    main_cr = nd('Div', [nd('ReduceSum', [nd('Mul', [main_rp, 'row_1d'], [[1, 1, 10, 1]]), 'axes2'], [[1, 1, 1, 1]], keepdims=1), main_total], [[1, 1, 1, 1]])
    # Main centroid col = sum(main_cp * col_1d) / main_total
    main_cc = nd('Div', [nd('ReduceSum', [nd('Mul', [main_cp, 'col_1d'], [[1, 1, 1, 10]]), 'axes3'], [[1, 1, 1, 1]], keepdims=1), main_total], [[1, 1, 1, 1]])

    axis_total = nd('ReduceSum', [axis_rp, 'axes2'], [[1, 1, 1, 1]], keepdims=1)
    axis_cr = nd('Div', [nd('ReduceSum', [nd('Mul', [axis_rp, 'row_1d'], [[1, 1, 10, 1]]), 'axes2'], [[1, 1, 1, 1]], keepdims=1), axis_total], [[1, 1, 1, 1]])
    axis_cc = nd('Div', [nd('ReduceSum', [nd('Mul', [axis_cp, 'col_1d'], [[1, 1, 1, 10]]), 'axes3'], [[1, 1, 1, 1]], keepdims=1), axis_total], [[1, 1, 1, 1]])

    # === STEP 6: Orientation ===
    dr = nd('Sub', [axis_cr, main_cr], [[1, 1, 1, 1]])
    dc = nd('Sub', [axis_cc, main_cc], [[1, 1, 1, 1]])
    abs_dr = nd('Abs', [dr], [[1, 1, 1, 1]])
    abs_dc = nd('Abs', [dc], [[1, 1, 1, 1]])
    is_horiz_b = nd('Greater', [abs_dr, nd('Sub', [abs_dc, 'c_half'], [[1, 1, 1, 1]])],
                    [([1, 1, 1, 1], TensorProto.BOOL)])
    is_horiz = nd('Cast', [is_horiz_b], [[1, 1, 1, 1]], to=1)
    is_vert = nd('Sub', ['c_one', is_horiz], [[1, 1, 1, 1]])

    # === STEP 7: Edges from 1D projections ===
    # has_row[r] = main_rp[r] > 0 → binary [1,1,10,1]
    has_mr_b = nd('Greater', [main_rp, 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
    has_mr = nd('Cast', [has_mr_b], [[1, 1, 10, 1]], to=1)
    no_mr = nd('Sub', ['c_one', has_mr], [[1, 1, 10, 1]])
    # Max main row: masked ArgMax
    mr_for_max = nd('Add', [nd('Mul', ['row_1d', has_mr], [[1, 1, 10, 1]]),
                             nd('Mul', ['c_neg_big', no_mr], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
    max_mr = nd('ReduceMax', [mr_for_max, 'axes2'], [[1, 1, 1, 1]], keepdims=1)
    # Min main row
    mr_for_min = nd('Add', [nd('Mul', ['row_1d', has_mr], [[1, 1, 10, 1]]),
                             nd('Mul', ['c_big', no_mr], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
    min_mr = nd('ReduceMin', [mr_for_min, 'axes2'], [[1, 1, 1, 1]], keepdims=1)

    # Max/min main col from col projection
    has_mc_b = nd('Greater', [main_cp, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)])
    has_mc = nd('Cast', [has_mc_b], [[1, 1, 1, 10]], to=1)
    no_mc = nd('Sub', ['c_one', has_mc], [[1, 1, 1, 10]])
    mc_for_max = nd('Add', [nd('Mul', ['col_1d', has_mc], [[1, 1, 1, 10]]),
                             nd('Mul', ['c_neg_big', no_mc], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    max_mc = nd('ReduceMax', [mc_for_max, 'axes3'], [[1, 1, 1, 1]], keepdims=1)
    mc_for_min = nd('Add', [nd('Mul', ['col_1d', has_mc], [[1, 1, 1, 10]]),
                             nd('Mul', ['c_big', no_mc], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    min_mc = nd('ReduceMin', [mc_for_min, 'axes3'], [[1, 1, 1, 1]], keepdims=1)

    # Axis edges from axis projections
    has_ar_b = nd('Greater', [axis_rp, 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
    has_ar = nd('Cast', [has_ar_b], [[1, 1, 10, 1]], to=1)
    no_ar = nd('Sub', ['c_one', has_ar], [[1, 1, 10, 1]])
    ar_for_max = nd('Add', [nd('Mul', ['row_1d', has_ar], [[1, 1, 10, 1]]),
                             nd('Mul', ['c_neg_big', no_ar], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
    max_ar = nd('ReduceMax', [ar_for_max, 'axes2'], [[1, 1, 1, 1]], keepdims=1)
    ar_for_min = nd('Add', [nd('Mul', ['row_1d', has_ar], [[1, 1, 10, 1]]),
                             nd('Mul', ['c_big', no_ar], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
    min_ar = nd('ReduceMin', [ar_for_min, 'axes2'], [[1, 1, 1, 1]], keepdims=1)

    has_ac_b = nd('Greater', [axis_cp, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)])
    has_ac = nd('Cast', [has_ac_b], [[1, 1, 1, 10]], to=1)
    no_ac = nd('Sub', ['c_one', has_ac], [[1, 1, 1, 10]])
    ac_for_max = nd('Add', [nd('Mul', ['col_1d', has_ac], [[1, 1, 1, 10]]),
                             nd('Mul', ['c_neg_big', no_ac], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    max_ac = nd('ReduceMax', [ac_for_max, 'axes3'], [[1, 1, 1, 1]], keepdims=1)
    ac_for_min = nd('Add', [nd('Mul', ['col_1d', has_ac], [[1, 1, 1, 10]]),
                             nd('Mul', ['c_big', no_ac], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    min_ac = nd('ReduceMin', [ac_for_min, 'axes3'], [[1, 1, 1, 1]], keepdims=1)

    # === STEP 8: Mirror position ===
    dr_pos_b = nd('Greater', [dr, 'c_zero'], [([1, 1, 1, 1], TensorProto.BOOL)])
    dr_pos = nd('Cast', [dr_pos_b], [[1, 1, 1, 1]], to=1)
    dr_neg = nd('Sub', ['c_one', dr_pos], [[1, 1, 1, 1]])
    h_sum = nd('Add', [nd('Add', [nd('Mul', [max_mr, dr_pos], [[1, 1, 1, 1]]),
                                    nd('Mul', [min_mr, dr_neg], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]),
                        nd('Add', [nd('Mul', [min_ar, dr_pos], [[1, 1, 1, 1]]),
                                    nd('Mul', [max_ar, dr_neg], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
    h_mirror = nd('Div', [h_sum, 'c_two'], [[1, 1, 1, 1]])

    dc_pos_b = nd('Greater', [dc, 'c_zero'], [([1, 1, 1, 1], TensorProto.BOOL)])
    dc_pos = nd('Cast', [dc_pos_b], [[1, 1, 1, 1]], to=1)
    dc_neg = nd('Sub', ['c_one', dc_pos], [[1, 1, 1, 1]])
    v_sum = nd('Add', [nd('Add', [nd('Mul', [max_mc, dc_pos], [[1, 1, 1, 1]]),
                                    nd('Mul', [min_mc, dc_neg], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]),
                        nd('Add', [nd('Mul', [min_ac, dc_pos], [[1, 1, 1, 1]]),
                                    nd('Mul', [max_ac, dc_neg], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
    v_mirror = nd('Div', [v_sum, 'c_two'], [[1, 1, 1, 1]])

    # === STEP 9: Reflected coords & Gather ===
    # Horizontal: reflected row
    h_mx2 = nd('Mul', [h_mirror, 'c_two'], [[1, 1, 1, 1]])
    h_rr = nd('Floor', [nd('Add', [nd('Sub', [h_mx2, 'row_grid'], [[1, 1, 10, 1]]), 'c_half'], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
    h_rc = nd('Clip', [h_rr, 'c_0f', 'c_9_0'], [[1, 1, 10, 1]])
    h_idx_f = nd('Mul', [h_rc, 'ones_1_1_1_10'], [[1, 1, 10, 10]])
    h_idx = nd('Cast', [h_idx_f], [([1, 1, 10, 10], TensorProto.INT64)], to=7)
    h_refl = nd('GatherElements', [main_mask, h_idx], [[1, 1, 10, 10]], axis=2)

    # Vertical: reflected col
    v_mx2 = nd('Mul', [v_mirror, 'c_two'], [[1, 1, 1, 1]])
    v_rc = nd('Floor', [nd('Add', [nd('Sub', [v_mx2, 'col_grid'], [[1, 1, 1, 10]]), 'c_half'], [[1, 1, 1, 10]])], [[1, 1, 1, 10]])
    v_rclip = nd('Clip', [v_rc, 'c_0f', 'c_9_0'], [[1, 1, 1, 10]])
    v_idx_f = nd('Mul', [v_rclip, 'ones_1_1_10_1'], [[1, 1, 10, 10]])
    v_idx = nd('Cast', [v_idx_f], [([1, 1, 10, 10], TensorProto.INT64)], to=7)
    v_refl = nd('GatherElements', [main_mask, v_idx], [[1, 1, 10, 10]], axis=3)

    # === STEP 10: Select & combine ===
    refl_mask = nd('Add', [nd('Mul', [is_horiz, h_refl], [[1, 1, 10, 10]]),
                            nd('Mul', [is_vert, v_refl], [[1, 1, 10, 10]])], [[1, 1, 10, 10]])
    combined = nd('Max', [main_mask, refl_mask], [[1, 1, 10, 10]])

    # === STEP 11: Output ===
    main_out = nd('Mul', [mc_oh_4d, combined], [[1, 10, 10, 10]])
    bg_mask = nd('Sub', ['c_one', combined], [[1, 1, 10, 10]])
    c3_oh = nd('OneHot', ['idx_3', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
    c3_oh_4d = nd('Reshape', [c3_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
    bg_out = nd('Mul', [c3_oh_4d, bg_mask], [[1, 10, 10, 10]])
    out_10 = nd('Add', [main_out, bg_out], [[1, 10, 10, 10]])

    # === STEP 12: Pad ===
    final = nd('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], [[1, 10, 30, 30]])

    x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 10, 30, 30])
    y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 10, 30, 30])
    nodes.append(helper.make_node('Identity', [final], ['output']))

    graph = helper.make_graph(nodes, 'task062', [x], [y], initializer=inits, value_info=vis)
    model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)])
    return model


if __name__ == '__main__':
    import onnxruntime as ort
    import json, zipfile
    from huggingface_hub import hf_hub_download

    print("Building Task 062 ONNX model v3...")
    model = build_task062_v3()
    
    output_path = '/app/task062_v3.onnx'
    del model.graph.value_info[:]
    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
    onnx.save(model, output_path)
    
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  File size: {os.path.getsize(output_path):,} bytes")

    # Validate
    sess = ort.InferenceSession(output_path)
    task_path = hf_hub_download('rogermt/neurogolf-solver', 'own-solver/neurogolf-2026.zip')
    with zipfile.ZipFile(task_path, 'r') as zf:
        data = json.loads(zf.read('task062.json'))

    all_examples = data['train'] + data['test'] + data.get('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 c, v in enumerate(row):
                if r < 30 and c < 30: inp[0][v][r][c] = 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 c, v in enumerate(row):
                if r < 30 and c < 30: exp[0][v][r][c] = 1.0
        if np.array_equal(out, exp): right_count += 1
        else:
            wrong_count += 1
            if wrong_count <= 3:
                print(f"  FAIL {i}: {len(np.where(out != exp)[0])} diffs")

    print(f"  Results: {right_count}/{right_count+wrong_count} pass ({wrong_count} fail)")

    if wrong_count == 0:
        # Quick static estimate
        params = sum(int(np.prod(init.dims)) for init in model.graph.initializer)
        mem_est = sum(int(np.prod([d.dim_value for d in vi.type.tensor_type.shape.dim])) * 
                      np.dtype(onnx.helper.tensor_dtype_to_np_dtype(vi.type.tensor_type.elem_type)).itemsize
                      for vi in model.graph.value_info
                      if vi.type.HasField('tensor_type') and vi.type.tensor_type.HasField('shape'))
        score_est = max(1.0, 25.0 - math.log(max(1.0, mem_est + params)))
        print(f"  Params: {params:,}, Memory: {mem_est:,}, Score: {score_est:.3f}")
        print(f"  Gain vs base (~11.5): +{score_est - 11.5:.3f}")