"""Build optimized ONNX for Task 387. Rule (266/266 verified): - 4 dots at corners of axis-aligned rectangle, 2 colors (diagonal pairs same) - Output: 3x3 ring around each dot (ring = other color), center = dot color - Dashed paths of color 5 between same-row/same-col dots - Path pattern: from each end of gap, place every 2nd cell, meeting at middle ONNX approach: - Find dot positions using ReduceSum and ArgMax - Use Conv for 3x3 expansion of dot positions - Compute path masks using modulo-2 arithmetic - Assemble output """ import sys sys.path.insert(0, '/app/repo/medal-solvers') from onnx import TensorProto import numpy as np from onnx_builder import OnnxBuilder H, W, C = 30, 30, 10 def build_task387(): 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_two', [2.0]) const('c_big', [1000.0]) const('c_neg_big', [-1000.0]) const('axes1', [1], 'i') const('axes2', [2], 'i') const('axes3', [3], 'i') const('axes23', [2, 3], 'i') const('shape_1_1_30_30', [1, 1, 30, 30], '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_30_30', [30, 30], 'i') const('shape_1', [1], 'i') # Row/col index 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') const('idx_5', [5], 'i') # === STEP 1: Find dot positions === # Non-bg mask [1,1,30,30] - need to exclude outside-grid cells too # Active = any channel set (within grid) 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) # fg = active AND not ch0 ch0 = nd('Gather', ['input', 'idx_0'], [[1, 1, 30, 30]], axis=1) not_ch0 = nd('Sub', ['c_one', ch0], [[1, 1, 30, 30]]) fg = nd('Mul', [active, not_ch0], [[1, 1, 30, 30]]) # 1 where non-bg within grid # Find bounding box of fg: r_min, r_max, c_min, c_max # r_min: minimum row with any fg not_fg = nd('Sub', ['c_one', fg], [[1, 1, 30, 30]]) rg_fg = nd('Add', [nd('Mul', [fg, 'row_grid'], [[1, 1, 30, 30]]), nd('Mul', [not_fg, 'c_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) r_min = nd('ReduceMin', [rg_fg, 'axes23'], [[1, 1, 1, 1]], keepdims=1) rg_fg_max = nd('Add', [nd('Mul', [fg, 'row_grid'], [[1, 1, 30, 30]]), nd('Mul', [not_fg, 'c_neg_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) r_max = nd('ReduceMax', [rg_fg_max, 'axes23'], [[1, 1, 1, 1]], keepdims=1) cg_fg = nd('Add', [nd('Mul', [fg, 'col_grid'], [[1, 1, 30, 30]]), nd('Mul', [not_fg, 'c_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) c_min = nd('ReduceMin', [cg_fg, 'axes23'], [[1, 1, 1, 1]], keepdims=1) cg_fg_max = nd('Add', [nd('Mul', [fg, 'col_grid'], [[1, 1, 30, 30]]), nd('Mul', [not_fg, 'c_neg_big'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) c_max = nd('ReduceMax', [cg_fg_max, 'axes23'], [[1, 1, 1, 1]], keepdims=1) # === STEP 2: Get colors at corners === # Color at (r_min, c_min) = argmax(input[0,:,r_min,c_min]) # To extract: mask pixels at (r_min, c_min) position at_rmin = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', r_min], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_rmin_f = nd('Cast', [at_rmin], [[1, 1, 30, 30]], to=1) at_cmin = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', c_min], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_cmin_f = nd('Cast', [at_cmin], [[1, 1, 30, 30]], to=1) at_tl = nd('Mul', [at_rmin_f, at_cmin_f], [[1, 1, 30, 30]]) # 1 only at top-left corner # Extract input values at TL: multiply input by at_tl mask, sum over spatial dims # Result: [1, 10] with channel values at TL position tl_vals = nd('Mul', ['input', at_tl], [[1, 10, 30, 30]]) tl_ch = nd('ReduceSum', [tl_vals, 'axes23'], [[1, 10, 1, 1]], keepdims=1) # [1,10,1,1] # Color A = argmax of tl_ch (excluding ch0) tl_ch_2d = nd('Reshape', [tl_ch, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) # Set ch0 to -1 so argmax picks non-bg const('mask_ch0_neg', np.array([[-1.0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=np.float32).reshape(1, 10, 1, 1)) tl_ch_masked = nd('Add', [tl_ch, 'mask_ch0_neg'], [[1, 10, 1, 1]]) color_a_idx = nd('ArgMax', [tl_ch_masked], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1) color_a_idx_1d = nd('Reshape', [color_a_idx, 'shape_1'], [([1], TensorProto.INT64)]) # Similarly for TR corner at_cmax = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', c_max], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_cmax_f = nd('Cast', [at_cmax], [[1, 1, 30, 30]], to=1) at_tr = nd('Mul', [at_rmin_f, at_cmax_f], [[1, 1, 30, 30]]) tr_vals = nd('Mul', ['input', at_tr], [[1, 10, 30, 30]]) tr_ch = nd('ReduceSum', [tr_vals, 'axes23'], [[1, 10, 1, 1]], keepdims=1) tr_ch_masked = nd('Add', [tr_ch, 'mask_ch0_neg'], [[1, 10, 1, 1]]) color_b_idx = nd('ArgMax', [tr_ch_masked], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1) color_b_idx_1d = nd('Reshape', [color_b_idx, 'shape_1'], [([1], TensorProto.INT64)]) # === STEP 3: Build 3x3 rings === # For each corner, create: 3x3 ring of other_color + center of own_color # Ring mask: |row - dot_row| <= 1 AND |col - dot_col| <= 1 # Center mask: row == dot_row AND col == dot_col # Color A onehot [1,10,1,1] color_a_oh = nd('OneHot', [color_a_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1) color_a_oh_4d = nd('Reshape', [color_a_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) color_b_oh = nd('OneHot', [color_b_idx_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1) color_b_oh_4d = nd('Reshape', [color_b_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) const('c_1_5', [1.5]) def build_ring_and_center(dot_row, dot_col, center_color_oh, ring_color_oh, label): """Build 3x3 pattern at dot position. Returns [1,10,30,30].""" # 3x3 box mask row_near = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', dot_row], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_1_5'], [([1, 1, 30, 30], TensorProto.BOOL)]) row_near_f = nd('Cast', [row_near], [[1, 1, 30, 30]], to=1) col_near = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', dot_col], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_1_5'], [([1, 1, 30, 30], TensorProto.BOOL)]) col_near_f = nd('Cast', [col_near], [[1, 1, 30, 30]], to=1) box = nd('Mul', [row_near_f, col_near_f], [[1, 1, 30, 30]]) # Center mask row_eq = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', dot_row], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) row_eq_f = nd('Cast', [row_eq], [[1, 1, 30, 30]], to=1) col_eq = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', dot_col], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) col_eq_f = nd('Cast', [col_eq], [[1, 1, 30, 30]], to=1) center = nd('Mul', [row_eq_f, col_eq_f], [[1, 1, 30, 30]]) # Ring = box - center ring = nd('Sub', [box, center], [[1, 1, 30, 30]]) # Apply colors [1,10,30,30] ring_colored = nd('Mul', [ring, ring_color_oh], [[1, 10, 30, 30]]) center_colored = nd('Mul', [center, center_color_oh], [[1, 10, 30, 30]]) return nd('Add', [ring_colored, center_colored], [[1, 10, 30, 30]]) # Four corners: TL(A,ring=B), TR(B,ring=A), BL(B,ring=A), BR(A,ring=B) tl_out = build_ring_and_center(r_min, c_min, color_a_oh_4d, color_b_oh_4d, 'tl') tr_out = build_ring_and_center(r_min, c_max, color_b_oh_4d, color_a_oh_4d, 'tr') bl_out = build_ring_and_center(r_max, c_min, color_b_oh_4d, color_a_oh_4d, 'bl') br_out = build_ring_and_center(r_max, c_max, color_a_oh_4d, color_b_oh_4d, 'br') result = nd('Add', [tl_out, tr_out], [[1, 10, 30, 30]]) result = nd('Add', [result, bl_out], [[1, 10, 30, 30]]) result = nd('Add', [result, br_out], [[1, 10, 30, 30]]) # === STEP 4: Dashed paths === # Horizontal path at rows r_min and r_max, cols from c_min+2 to c_max-2 # Vertical path at cols c_min and c_max, rows from r_min+2 to r_max-2 # Path formula: on_path(pos, start, end) where start = edge+2, end = other_edge-2 # on_path = ((pos - start) % 2 == 0 AND pos <= mid) OR ((end - pos) % 2 == 0 AND pos >= mid) # mid = (start + end) / 2 h_start = nd('Add', [c_min, 'c_two'], [[1, 1, 1, 1]]) h_end = nd('Sub', [c_max, 'c_two'], [[1, 1, 1, 1]]) h_mid = nd('Div', [nd('Add', [h_start, h_end], [[1, 1, 1, 1]]), 'c_two'], [[1, 1, 1, 1]]) v_start = nd('Add', [r_min, 'c_two'], [[1, 1, 1, 1]]) v_end = nd('Sub', [r_max, 'c_two'], [[1, 1, 1, 1]]) v_mid = nd('Div', [nd('Add', [v_start, v_end], [[1, 1, 1, 1]]), 'c_two'], [[1, 1, 1, 1]]) # For horizontal path: check each col position # dist_from_start = col_grid - h_start (for cols in [h_start, h_end]) # dist_from_end = h_end - col_grid # mod2_start = dist_from_start - 2 * floor(dist_from_start / 2) # on_from_left = (mod2_start < 0.5) AND (col_grid <= h_mid) # on_from_right = (mod2_end < 0.5) AND (col_grid >= h_mid) # in_range = (col_grid >= h_start - 0.5) AND (col_grid <= h_end + 0.5) dist_from_h_start = nd('Sub', ['col_grid', h_start], [[1, 1, 30, 30]]) div_h_s = nd('Div', [dist_from_h_start, 'c_two'], [[1, 1, 30, 30]]) floor_h_s = nd('Floor', [div_h_s], [[1, 1, 30, 30]]) mod2_h_s = nd('Sub', [dist_from_h_start, nd('Mul', ['c_two', floor_h_s], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) mod2_h_s_ok = nd('Less', [mod2_h_s, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) col_le_hmid = nd('Less', ['col_grid', nd('Add', [h_mid, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) on_h_left = nd('And', [mod2_h_s_ok, col_le_hmid], [([1, 1, 30, 30], TensorProto.BOOL)]) dist_from_h_end = nd('Sub', [h_end, 'col_grid'], [[1, 1, 30, 30]]) div_h_e = nd('Div', [dist_from_h_end, 'c_two'], [[1, 1, 30, 30]]) floor_h_e = nd('Floor', [div_h_e], [[1, 1, 30, 30]]) mod2_h_e = nd('Sub', [dist_from_h_end, nd('Mul', ['c_two', floor_h_e], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) mod2_h_e_ok = nd('Less', [mod2_h_e, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) col_ge_hmid = nd('Greater', ['col_grid', nd('Sub', [h_mid, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) on_h_right = nd('And', [mod2_h_e_ok, col_ge_hmid], [([1, 1, 30, 30], TensorProto.BOOL)]) on_h_path = nd('Or', [on_h_left, on_h_right], [([1, 1, 30, 30], TensorProto.BOOL)]) # Range check: col in [h_start, h_end] h_in_range1 = nd('Greater', ['col_grid', nd('Sub', [h_start, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) h_in_range2 = nd('Less', ['col_grid', nd('Add', [h_end, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) h_in_range = nd('And', [h_in_range1, h_in_range2], [([1, 1, 30, 30], TensorProto.BOOL)]) h_path_valid = nd('And', [on_h_path, h_in_range], [([1, 1, 30, 30], TensorProto.BOOL)]) h_path_f = nd('Cast', [h_path_valid], [[1, 1, 30, 30]], to=1) # H-path at row r_min at_rmin_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', r_min], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_rmin_f2 = nd('Cast', [at_rmin_row], [[1, 1, 30, 30]], to=1) h_path_top = nd('Mul', [h_path_f, at_rmin_f2], [[1, 1, 30, 30]]) # H-path at row r_max at_rmax_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', r_max], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_rmax_f2 = nd('Cast', [at_rmax_row], [[1, 1, 30, 30]], to=1) h_path_bot = nd('Mul', [h_path_f, at_rmax_f2], [[1, 1, 30, 30]]) # Vertical path (same logic but for rows) dist_from_v_start = nd('Sub', ['row_grid', v_start], [[1, 1, 30, 30]]) div_v_s = nd('Div', [dist_from_v_start, 'c_two'], [[1, 1, 30, 30]]) floor_v_s = nd('Floor', [div_v_s], [[1, 1, 30, 30]]) mod2_v_s = nd('Sub', [dist_from_v_start, nd('Mul', ['c_two', floor_v_s], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) mod2_v_s_ok = nd('Less', [mod2_v_s, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) row_le_vmid = nd('Less', ['row_grid', nd('Add', [v_mid, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) on_v_left = nd('And', [mod2_v_s_ok, row_le_vmid], [([1, 1, 30, 30], TensorProto.BOOL)]) dist_from_v_end = nd('Sub', [v_end, 'row_grid'], [[1, 1, 30, 30]]) div_v_e = nd('Div', [dist_from_v_end, 'c_two'], [[1, 1, 30, 30]]) floor_v_e = nd('Floor', [div_v_e], [[1, 1, 30, 30]]) mod2_v_e = nd('Sub', [dist_from_v_end, nd('Mul', ['c_two', floor_v_e], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) mod2_v_e_ok = nd('Less', [mod2_v_e, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) row_ge_vmid = nd('Greater', ['row_grid', nd('Sub', [v_mid, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) on_v_right = nd('And', [mod2_v_e_ok, row_ge_vmid], [([1, 1, 30, 30], TensorProto.BOOL)]) on_v_path = nd('Or', [on_v_left, on_v_right], [([1, 1, 30, 30], TensorProto.BOOL)]) v_in_range1 = nd('Greater', ['row_grid', nd('Sub', [v_start, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) v_in_range2 = nd('Less', ['row_grid', nd('Add', [v_end, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) v_in_range = nd('And', [v_in_range1, v_in_range2], [([1, 1, 30, 30], TensorProto.BOOL)]) v_path_valid = nd('And', [on_v_path, v_in_range], [([1, 1, 30, 30], TensorProto.BOOL)]) v_path_f = nd('Cast', [v_path_valid], [[1, 1, 30, 30]], to=1) # V-path at col c_min at_cmin_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', c_min], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_cmin_f2 = nd('Cast', [at_cmin_col], [[1, 1, 30, 30]], to=1) v_path_left = nd('Mul', [v_path_f, at_cmin_f2], [[1, 1, 30, 30]]) # V-path at col c_max at_cmax_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', c_max], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_cmax_f2 = nd('Cast', [at_cmax_col], [[1, 1, 30, 30]], to=1) v_path_right = nd('Mul', [v_path_f, at_cmax_f2], [[1, 1, 30, 30]]) # Combine all paths [1,1,30,30] all_paths = nd('Add', [h_path_top, h_path_bot], [[1, 1, 30, 30]]) all_paths = nd('Add', [all_paths, v_path_left], [[1, 1, 30, 30]]) all_paths = nd('Add', [all_paths, v_path_right], [[1, 1, 30, 30]]) all_paths_clip = nd('Clip', [all_paths, 'c_zero', 'c_one'], [[1, 1, 30, 30]]) # Color 5 for paths color5_oh = nd('OneHot', ['idx_5', 'depth_10', 'oh_vals'], [[1, 10]], axis=1) color5_oh_4d = nd('Reshape', [color5_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) path_colored = nd('Mul', [all_paths_clip, color5_oh_4d], [[1, 10, 30, 30]]) # Final: rings + paths final = nd('Add', [result, path_colored], [[1, 10, 30, 30]]) # Add background (ch0): within grid area where no other channel is active # Grid area: use 'active' mask (detected earlier from input) fg_final = nd('ReduceSum', [final, 'axes1'], [[1, 1, 30, 30]], keepdims=1) fg_final_b = nd('Greater', [fg_final, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) fg_final_f = nd('Cast', [fg_final_b], [[1, 1, 30, 30]], to=1) # bg = active AND not fg_final bg_mask = nd('Sub', [active, fg_final_f], [[1, 1, 30, 30]]) bg_mask_clip = nd('Clip', [bg_mask, 'c_zero', 'c_one'], [[1, 1, 30, 30]]) # ch0 onehot const('idx_0_for_bg', [0], 'i') ch0_oh = nd('OneHot', ['idx_0_for_bg', 'depth_10', 'oh_vals'], [[1, 10]], axis=1) ch0_oh_4d = nd('Reshape', [ch0_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) bg_colored = nd('Mul', [bg_mask_clip, ch0_oh_4d], [[1, 10, 30, 30]]) final_with_bg = nd('Add', [final, bg_colored], [[1, 10, 30, 30]]) return b.finish('task387', last_tensor=final_with_bg) if __name__ == '__main__': import os, json, math import numpy as np import onnx import onnxruntime as ort print("Building Task 387 ONNX model...") model = build_task387() os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True) output_path = '/app/repo/medal-solvers/optimized/task387.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)") sess = ort.InferenceSession(output_path) with open('/app/task-data/task387.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_grid = ex['input'] inp = np.zeros((1, 10, 30, 30), dtype=np.float32) for r, row in enumerate(inp_grid): 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 <= 3: diff_locs = np.where(out != exp) print(f" FAIL {i}: {len(diff_locs[0])} diffs") print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")