"""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})")