"""Build optimized ONNX model for Task 025. Algorithm (266/266 verified): - Detect h/v lines for each color - For each dot of color c, project it adjacent to nearest same-color line - Only one line per color (never both h and v for same color) Score target: ~11-12 pts (from base 8.71) """ 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_task025(): b = OnnxBuilder() const, nd = b.const, b.nd # === CONSTANTS === const('c_half', [0.5]) const('c_one', [1.0]) const('c_zero', [0.0]) const('axes1', [1], 'i') const('axes2', [2], 'i') const('axes3', [3], 'i') const('shape_1_1_30_30', [1, 1, 30, 30], 'i') const('shape_1_1_30_1', [1, 1, 30, 1], 'i') const('shape_1_1_1_30', [1, 1, 1, 30], 'i') const('shape_1_1_1_1', [1, 1, 1, 1], 'i') const('shape_1_10_1_1', [1, 10, 1, 1], 'i') # Row/col index grids as [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]) for i in range(10): const(f'idx_{i}', [i], 'i') # === STEP 1: Grid active mask [1,1,30,30] === 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) # Grid width per row [1,1,30,1] grid_w = nd('ReduceSum', [active, 'axes3'], [[1, 1, 30, 1]], keepdims=1) # Grid height per col [1,1,1,30] grid_h = nd('ReduceSum', [active, 'axes2'], [[1, 1, 1, 30]], keepdims=1) # === STEP 2: Process each color === const('zeros_1_10_30_30', np.zeros((1, 10, 30, 30), dtype=np.float32)) delta = 'zeros_1_10_30_30' for c in range(1, 10): # Channel c mask [1,1,30,30] ch = nd('Gather', ['input', f'idx_{c}'], [[1, 1, 30, 30]], axis=1) # Row sum of color c [1,1,30,1] ch_row_sum = nd('ReduceSum', [ch, 'axes3'], [[1, 1, 30, 1]], keepdims=1) # Col sum of color c [1,1,1,30] ch_col_sum = nd('ReduceSum', [ch, 'axes2'], [[1, 1, 1, 30]], keepdims=1) # H-line: row r is a line if ch_row_sum[r] == grid_w[r] > 0 diff_h = nd('Sub', [ch_row_sum, grid_w], [[1, 1, 30, 1]]) abs_diff_h = nd('Abs', [diff_h], [[1, 1, 30, 1]]) h_match_b = nd('Less', [abs_diff_h, 'c_half'], [([1, 1, 30, 1], TensorProto.BOOL)]) gw_pos_b = nd('Greater', [grid_w, 'c_half'], [([1, 1, 30, 1], TensorProto.BOOL)]) h_line_b = nd('And', [h_match_b, gw_pos_b], [([1, 1, 30, 1], TensorProto.BOOL)]) h_line = nd('Cast', [h_line_b], [[1, 1, 30, 1]], to=1) # [1,1,30,1] # V-line: col is a line if ch_col_sum[col] == grid_h[col] > 0 diff_v = nd('Sub', [ch_col_sum, grid_h], [[1, 1, 1, 30]]) abs_diff_v = nd('Abs', [diff_v], [[1, 1, 1, 30]]) v_match_b = nd('Less', [abs_diff_v, 'c_half'], [([1, 1, 1, 30], TensorProto.BOOL)]) gh_pos_b = nd('Greater', [grid_h, 'c_half'], [([1, 1, 1, 30], TensorProto.BOOL)]) v_line_b = nd('And', [v_match_b, gh_pos_b], [([1, 1, 1, 30], TensorProto.BOOL)]) v_line = nd('Cast', [v_line_b], [[1, 1, 1, 30]], to=1) # [1,1,1,30] # has_h_line / has_v_line (scalar [1,1,1,1]) has_h = nd('ReduceMax', [h_line, 'axes2'], [[1, 1, 1, 1]], keepdims=1) has_v = nd('ReduceMax', [v_line, 'axes3'], [[1, 1, 1, 1]], keepdims=1) # H-line position: weighted sum of h_line * row_grid[...,0] # h_line is [1,1,30,1], row_grid is [1,1,30,30] # Multiply and sum over row dim (axis 2) # Use row values directly: [1,1,30,1] * row_indices [1,1,30,1] const(f'row_idx_col_{c}', np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1)) h_pos = nd('ReduceSum', [nd('Mul', [h_line, f'row_idx_col_{c}'], [[1, 1, 30, 1]]), 'axes2'], [[1, 1, 1, 1]], keepdims=1) # [1,1,1,1] = line row position # V-line position const(f'col_idx_row_{c}', np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30)) v_pos = nd('ReduceSum', [nd('Mul', [v_line, f'col_idx_row_{c}'], [[1, 1, 1, 30]]), 'axes3'], [[1, 1, 1, 1]], keepdims=1) # [1,1,1,1] = line col position # On-line mask [1,1,30,30]: broadcast h_line [1,1,30,1] with v_line [1,1,1,30] on_line_mask = nd('Max', [h_line, v_line], [[1, 1, 30, 30]]) # broadcasts correctly! on_line_clip = nd('Clip', [on_line_mask, 'c_zero', 'c_one'], [[1, 1, 30, 30]]) not_on_line = nd('Sub', ['c_one', on_line_clip], [[1, 1, 30, 30]]) # Dots = color c pixels not on a line dots = nd('Mul', [ch, not_on_line], [[1, 1, 30, 30]]) # === H-line projection === # Dots above line: row < h_pos above_b = nd('Less', ['row_grid', nd('Sub', [h_pos, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) above_f = nd('Cast', [above_b], [[1, 1, 30, 30]], to=1) below_f = nd('Sub', ['c_one', above_f], [[1, 1, 30, 30]]) # Collapse dots above/below to column indicators dots_above = nd('Mul', [dots, above_f], [[1, 1, 30, 30]]) dots_below = nd('Mul', [dots, below_f], [[1, 1, 30, 30]]) # Which cols have dots above/below? cols_above = nd('ReduceMax', [dots_above, 'axes2'], [[1, 1, 1, 30]], keepdims=1) # [1,1,1,30] cols_below = nd('ReduceMax', [dots_below, 'axes2'], [[1, 1, 1, 30]], keepdims=1) # Target rows h_target_above = nd('Sub', [h_pos, 'c_one'], [[1, 1, 1, 1]]) # row above line h_target_below = nd('Add', [h_pos, 'c_one'], [[1, 1, 1, 1]]) # row below line # Masks for target rows [1,1,30,30] at_above_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', h_target_above], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_above_row_f = nd('Cast', [at_above_row], [[1, 1, 30, 30]], to=1) at_below_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', h_target_below], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_below_row_f = nd('Cast', [at_below_row], [[1, 1, 30, 30]], to=1) # H-projection: place at target row for cols with dots h_proj_above = nd('Mul', [at_above_row_f, cols_above], [[1, 1, 30, 30]]) h_proj_below = nd('Mul', [at_below_row_f, cols_below], [[1, 1, 30, 30]]) h_proj = nd('Add', [h_proj_above, h_proj_below], [[1, 1, 30, 30]]) h_proj_gated = nd('Mul', [h_proj, has_h], [[1, 1, 30, 30]]) # === V-line projection === left_b = nd('Less', ['col_grid', nd('Sub', [v_pos, 'c_half'], [[1, 1, 1, 1]])], [([1, 1, 30, 30], TensorProto.BOOL)]) left_f = nd('Cast', [left_b], [[1, 1, 30, 30]], to=1) right_f = nd('Sub', ['c_one', left_f], [[1, 1, 30, 30]]) dots_left = nd('Mul', [dots, left_f], [[1, 1, 30, 30]]) dots_right = nd('Mul', [dots, right_f], [[1, 1, 30, 30]]) rows_left = nd('ReduceMax', [dots_left, 'axes3'], [[1, 1, 30, 1]], keepdims=1) rows_right = nd('ReduceMax', [dots_right, 'axes3'], [[1, 1, 30, 1]], keepdims=1) v_target_left = nd('Sub', [v_pos, 'c_one'], [[1, 1, 1, 1]]) v_target_right = nd('Add', [v_pos, 'c_one'], [[1, 1, 1, 1]]) at_left_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', v_target_left], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_left_col_f = nd('Cast', [at_left_col], [[1, 1, 30, 30]], to=1) at_right_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', v_target_right], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) at_right_col_f = nd('Cast', [at_right_col], [[1, 1, 30, 30]], to=1) v_proj_left = nd('Mul', [at_left_col_f, rows_left], [[1, 1, 30, 30]]) v_proj_right = nd('Mul', [at_right_col_f, rows_right], [[1, 1, 30, 30]]) v_proj = nd('Add', [v_proj_left, v_proj_right], [[1, 1, 30, 30]]) v_proj_gated = nd('Mul', [v_proj, has_v], [[1, 1, 30, 30]]) # Combined projection proj = nd('Add', [h_proj_gated, v_proj_gated], [[1, 1, 30, 30]]) # Net change for this color: ALL dots are removed, proj added only if has_line # dots are always removed (they're not in the output regardless of whether there's a line) net_change = nd('Sub', [proj, dots], [[1, 1, 30, 30]]) # Expand to 10 channels c_oh = nd('OneHot', [f'idx_{c}', 'depth_10', 'oh_vals'], [[1, 10]], axis=1) c_oh_4d = nd('Reshape', [c_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) ch0_oh = nd('OneHot', ['idx_0', 'depth_10', 'oh_vals'], [[1, 10]], axis=1) ch0_oh_4d = nd('Reshape', [ch0_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) # delta_c = net_change * (c_onehot - ch0_onehot) oh_diff = nd('Sub', [c_oh_4d, ch0_oh_4d], [[1, 10, 1, 1]]) color_delta = nd('Mul', [net_change, oh_diff], [[1, 10, 30, 30]]) delta = nd('Add', [delta, color_delta], [[1, 10, 30, 30]]) # Output = input + delta, clipped to [0,1] final = nd('Add', ['input', delta], [[1, 10, 30, 30]]) final_clipped = nd('Clip', [final, 'c_zero', 'c_one'], [[1, 10, 30, 30]]) return b.finish('task025', last_tensor=final_clipped) if __name__ == '__main__': import os, json, math import numpy as np import onnx import onnxruntime as ort print("Building Task 025 ONNX model...") model = build_task025() os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True) output_path = '/app/repo/medal-solvers/optimized/task025.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)") # Validate sess = ort.InferenceSession(output_path) with open('/app/task-data/task025.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 <= 5: diff_locs = np.where(out != exp) print(f" FAIL {i}: {len(diff_locs[0])} diffs") for d in range(min(3, len(diff_locs[0]))): ch_d, r_d, c_d = diff_locs[1][d], diff_locs[2][d], diff_locs[3][d] print(f" ch={ch_d} ({r_d},{c_d}): pred={out[0,ch_d,r_d,c_d]} exp={exp[0,ch_d,r_d,c_d]}") print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}") if wrong_count == 0: 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])) * 4 for vi in model.graph.value_info if vi.type.tensor_type.shape.dim) score_est = max(1.0, 25.0 - math.log(max(1.0, mem_est + params))) print(f" Params: {params:,}, Memory est: {mem_est:,}, Score est: {score_est:.3f}")