"""Build optimized ONNX model for Task 200. Rule (84/84 verified): 1. Input: 10x10 grid with single colored dot at row 9 2. Output: From dot's column rightward, alternating pattern: - Even offsets (0,2,4,...): full vertical line of dot color - Odd offsets (1,3,5,...): color 5 at alternating top(row0)/bottom(row9) ONNX approach: - Slice to 10x10 - Detect dot column and color - Use column grid with modular arithmetic to create masks - Apply colors - Pad back to 30x30 Base: 91 nodes, 198010 params, score 10.33 Target: ~30 nodes, ~500 params → score ~15.5 (gain +5.2!) """ import sys, os import numpy as np import onnx from onnx import helper, numpy_helper, TensorProto import math def build_task200(): 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_two', [2.0]) const('c_four', [4.0]) const('c_1_5', [1.5]) const('c_8_5', [8.5]) 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('axes23', [2, 3], 'i') const('axes2', [2], 'i') # Column grid [1,1,1,10] const('col_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10)) # Row grid [1,1,10,1] const('row_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1)) 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') const('depth_10', [10.0]) const('oh_vals', [0.0, 1.0]) # Pad to 30x30 const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i') const('pad_val_zero', [0.0]) # Channel slice for non-bg const('sl_ch_start', [0, 1, 0, 0], 'i') const('sl_ch_end', [1, 10, 10, 10], 'i') # bg mask for channel selection (exclude ch0) const('bg_mask_10', np.array([[[[0, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]], dtype=np.float32).reshape(1, 10, 1, 1)) # Color 5 index const('idx_5', [5], 'i') const('idx_1_arr', [1], 'i') # === STEP 1: Slice input to 10x10 === inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]]) # === STEP 2: Find dot column and color === # Sum over channels to get active mask, then find the dot nonbg = nd('Slice', [inp10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]]) active = nd('ReduceSum', [nonbg, 'axes1'], [[1, 1, 10, 10]], keepdims=1) # Sum over rows to get column indicator [1,1,1,10] col_sums = nd('ReduceSum', [active, 'axes2'], [[1, 1, 1, 10]], keepdims=1) # Dot column = where col_sums > 0 → use col_grid * col_sums → max gives dot col col_weighted = nd('Mul', ['col_grid', col_sums], [[1, 1, 1, 10]]) dot_col_f = nd('ReduceMax', [col_weighted], [[1, 1, 1, 1]], keepdims=1) # [1,1,1,1] # Dot color: which channel has the dot? ch_sums = nd('ReduceSum', [inp10, 'axes23'], [[1, 10, 1, 1]], keepdims=1) ch_sums_nobg = nd('Mul', [ch_sums, 'bg_mask_10'], [[1, 10, 1, 1]]) dot_color_idx = nd('ArgMax', [ch_sums_nobg], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1) dot_color_1d = nd('Reshape', [dot_color_idx, 'shape_1'], [([1], TensorProto.INT64)]) # === STEP 3: Compute column masks === # relative column offset from dot: col - dot_col rel_col = nd('Sub', ['col_grid', dot_col_f], [[1, 1, 1, 10]]) # Only columns >= dot_col (rel >= 0) in_range_b = nd('Greater', [rel_col, nd('Sub', ['c_half', 'c_one'], [[1, 1, 1, 1]])], [([1, 1, 1, 10], TensorProto.BOOL)]) # rel >= 0 in_range = nd('Cast', [in_range_b], [[1, 1, 1, 10]], to=1) # Even offset columns (rel % 2 == 0): vertical lines # rel / 2, floor, * 2, compare to rel rel_div2 = nd('Div', [rel_col, 'c_two'], [[1, 1, 1, 10]]) rel_floor = nd('Floor', [rel_div2], [[1, 1, 1, 10]]) rel_x2 = nd('Mul', [rel_floor, 'c_two'], [[1, 1, 1, 10]]) is_even_diff = nd('Abs', [nd('Sub', [rel_col, rel_x2], [[1, 1, 1, 10]])], [[1, 1, 1, 10]]) is_even_b = nd('Less', [is_even_diff, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)]) is_even = nd('Cast', [is_even_b], [[1, 1, 1, 10]], to=1) # Vertical line mask: is_even AND in_range → [1,1,1,10] broadcast to [1,1,10,10] vert_mask_col = nd('Mul', [is_even, in_range], [[1, 1, 1, 10]]) # [1,1,1,10] # This broadcasts to all rows (since it's independent of row) # Odd offset columns: connectors is_odd = nd('Sub', ['c_one', is_even], [[1, 1, 1, 10]]) odd_mask_col = nd('Mul', [is_odd, in_range], [[1, 1, 1, 10]]) # Connector position: for the k-th odd column (k=0,1,2,...): # k even → row 0, k odd → row 9 # k = (rel - 1) / 2 for odd rel # k % 2 == 0 → top, k % 2 == 1 → bottom rel_minus1 = nd('Sub', [rel_col, 'c_one'], [[1, 1, 1, 10]]) k_val = nd('Div', [rel_minus1, 'c_two'], [[1, 1, 1, 10]]) # k = (rel-1)/2 for odd cols k_div2 = nd('Div', [k_val, 'c_two'], [[1, 1, 1, 10]]) k_floor = nd('Floor', [k_div2], [[1, 1, 1, 10]]) k_x2 = nd('Mul', [k_floor, 'c_two'], [[1, 1, 1, 10]]) k_mod2_diff = nd('Abs', [nd('Sub', [k_val, k_x2], [[1, 1, 1, 10]])], [[1, 1, 1, 10]]) k_is_even_b = nd('Less', [k_mod2_diff, 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)]) k_is_even = nd('Cast', [k_is_even_b], [[1, 1, 1, 10]], to=1) # 1 if top connector k_is_odd = nd('Sub', ['c_one', k_is_even], [[1, 1, 1, 10]]) # 1 if bottom connector # Top connector: odd_mask_col AND k_is_even AND row==0 # Bottom connector: odd_mask_col AND k_is_odd AND row==9 at_row0_b = nd('Less', ['row_grid', 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)]) at_row0 = nd('Cast', [at_row0_b], [[1, 1, 10, 1]], to=1) at_row9_b = nd('Greater', ['row_grid', 'c_8_5'], [([1, 1, 10, 1], TensorProto.BOOL)]) at_row9 = nd('Cast', [at_row9_b], [[1, 1, 10, 1]], to=1) # Top connector mask [1,1,10,10]: odd_mask_col * k_is_even * at_row0 top_conn = nd('Mul', [nd('Mul', [odd_mask_col, k_is_even], [[1, 1, 1, 10]]), at_row0], [[1, 1, 10, 10]]) # Bottom connector mask bot_conn = nd('Mul', [nd('Mul', [odd_mask_col, k_is_odd], [[1, 1, 1, 10]]), at_row9], [[1, 1, 10, 10]]) # Combined connector mask (color 5) conn_mask = nd('Add', [top_conn, bot_conn], [[1, 1, 10, 10]]) # === STEP 4: Apply colors === # Dot color one-hot dot_oh = nd('OneHot', [dot_color_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1) dot_oh_4d = nd('Reshape', [dot_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) # Vertical lines: dot_color * vert_mask vert_out = nd('Mul', [dot_oh_4d, vert_mask_col], [[1, 10, 10, 10]]) # Connectors: color 5 * conn_mask c5_oh = nd('OneHot', ['idx_5', 'depth_10', 'oh_vals'], [[1, 10]], axis=1) c5_oh_4d = nd('Reshape', [c5_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]]) conn_out = nd('Mul', [c5_oh_4d, conn_mask], [[1, 10, 10, 10]]) # Combine colored_10 = nd('Add', [vert_out, conn_out], [[1, 10, 10, 10]]) # === STEP 5: Add background channel === nonbg_out = nd('Slice', [colored_10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]]) any_color = nd('ReduceMax', [nonbg_out, 'axes1'], [[1, 1, 10, 10]], keepdims=1) bg_ch = nd('Sub', ['c_one', any_color], [[1, 1, 10, 10]]) const('ch0_idx', [0], 'i') ch0_oh = nd('OneHot', ['ch0_idx', '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_full = nd('Mul', [ch0_oh_4d, bg_ch], [[1, 10, 10, 10]]) out_10 = nd('Add', [colored_10, bg_full], [[1, 10, 10, 10]]) # === STEP 6: Pad to 30x30 === final = nd('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], [[1, 10, 30, 30]]) # Build model 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, 'task200', [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 200 ONNX model...") model = build_task200() output_path = '/app/task200.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") # 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('task200.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: 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)}") 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.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 est: {mem_est:,}") print(f" Score est: {score_est:.3f} (base: 10.326)") print(f" Gain est: +{score_est - 10.326:.3f}") print(f" ** THIS IS A STRONG CANDIDATE FOR SUBMISSION **")