""" Build optimized ONNX model for Task 285 using ScatterElements approach. Key insight: compute everything at ACTIVE pixels (where bbox is correct), then SCATTER results to destination positions. No need for bbox at destinations. Algorithm: 1. Flood fill (MaxPool 3x3, 30 iter) for component labels 2. Adjacency [900x900] for component membership 3. Dominant detection via MatMul + ArgMax 4. Bbox at active pixels via MaxPool within active mask 5. For each color c: find minority_c position per component via adjacency 6. For each dominant pixel: compute destination index in minority_c's cell 7. ScatterElements to write color at destination 8. Overlay with original input Results (validated with official neurogolf_utils.py): ARC-AGI: 4 pass, 0 fail ARC-GEN: 261 pass, 0 fail Score: 8.818 points (original: 5.98, gain: +2.84) Memory: 10,611,900 bytes (original: 181M, 17x reduction) Usage: # From medal-solvers/ directory: python build_task285_scatter.py \\ --neurogolf-utils ../own-solver/neurogolf_utils.py \\ --task-data-dir ../task-data \\ --output optimized/task285.onnx # Then create submission: python swap_and_submit.py \\ --base ../submission-6043.zip \\ --models optimized/task285.onnx \\ --output ../submission.zip Requirements: pip install onnx onnxruntime numpy matplotlib onnx_tool # Task data: extract own-solver/neurogolf-2026.zip → task-data/ """ import onnx from onnx import helper, TensorProto, numpy_helper import numpy as np import os H, W, C, N = 30, 30, 10, 900 def build_task285(): 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 def prop_masked(init_val, mask, iters=30): v = init_val for _ in range(iters): mp = nd('MaxPool', [v], [[1,1,H,W]], kernel_shape=[3,3], pads=[1,1,1,1]) v = nd('Mul', [mp, mask], [[1,1,H,W]]) return v # --- Constants --- const('c_half', [0.5]); const('c_one', [1.0]); const('c_two', [2.0]) const('c_zero', [0.0]); const('c_nm1', [float(N-1)]) const('c_hm1', [float(H-1)]); const('c_wm1', [float(W-1)]); const('c_w', [float(W)]) const('shape_n1', [N, 1], 'i'); const('shape_1n', [1, N], 'i') const('shape_n', [N], 'i'); const('shape_cn', [C, N], 'i') const('shape_hw', [1, 1, H, W], 'i'); const('shape_1chw', [1, C, H, W], 'i') const('axes1', [1], 'i') const('coord_labels', np.arange(1, N+1, dtype=np.float32).reshape(1,1,H,W)) row_s = np.arange(H, dtype=np.float32).reshape(1,1,H,1) * np.ones((1,1,1,W), dtype=np.float32) col_s = np.arange(W, dtype=np.float32).reshape(1,1,1,W) * np.ones((1,1,H,1), dtype=np.float32) const('row_s', row_s); const('col_s', col_s) const('row_flat', row_s.reshape(N)); const('col_flat', col_s.reshape(N)) const('bg_mask', np.array([0.0]+[1.0]*9, dtype=np.float32).reshape(1, C)) const('s1', [0,1,0,0], 'i'); const('e1', [1,C,H,W], 'i'); const('ax4', [0,1,2,3], 'i') const('zeros_n', np.zeros(N, dtype=np.float32)) const('ones_n', np.ones(N, dtype=np.float32)) for c in range(1, 10): const(f's_ch{c}', [0, c, 0, 0], 'i') const(f'e_ch{c}', [1, c+1, H, W], 'i') # STEP 1: Active mask ch19 = nd('Slice', ['input', 's1', 'e1', 'ax4'], [[1, 9, H, W]]) asum = nd('ReduceSum', [ch19, 'axes1'], [[1, 1, H, W]], keepdims=1) abool = nd('Greater', [asum, 'c_half'], [([1,1,H,W], TensorProto.BOOL)]) active = nd('Cast', [abool], [[1, 1, H, W]], to=1) # STEP 2: Flood fill (8-connected via MaxPool 3x3) label = nd('Mul', [active, 'coord_labels'], [[1,1,H,W]]) for _ in range(30): mp = nd('MaxPool', [label], [[1,1,H,W]], kernel_shape=[3,3], pads=[1,1,1,1]) label = nd('Mul', [mp, active], [[1,1,H,W]]) # STEP 3: Adjacency + dominant detection l_n1 = nd('Reshape', [label, 'shape_n1'], [[N, 1]]) l_1n = nd('Reshape', [label, 'shape_1n'], [[1, N]]) adj_b = nd('Equal', [l_n1, l_1n], [([N, N], TensorProto.BOOL)]) adj = nd('Cast', [adj_b], [[N, N]], to=1) a_flat = nd('Reshape', [active, 'shape_n'], [[N]]) a_col = nd('Reshape', [a_flat, 'shape_n1'], [[N, 1]]) adj_m = nd('Mul', [adj, a_col], [[N, N]]) inp_cn = nd('Reshape', ['input', 'shape_cn'], [[C, N]]) inp_nc = nd('Transpose', [inp_cn], [[N, C]], perm=[1, 0]) csums = nd('MatMul', [adj_m, inp_nc], [[N, C]]) csums_masked = nd('Mul', [csums, 'bg_mask'], [[N, C]]) dom_idx = nd('ArgMax', [csums_masked], [([N, 1], TensorProto.INT64)], axis=1, keepdims=1) pixel_color = nd('ArgMax', [inp_nc], [([N, 1], TensorProto.INT64)], axis=1, keepdims=1) is_dom_b = nd('Equal', [dom_idx, pixel_color], [([N, 1], TensorProto.BOOL)]) is_dom_f = nd('Cast', [is_dom_b], [[N, 1]], to=1) is_dom = nd('Reshape', [is_dom_f, 'shape_n'], [[N]]) is_dom_active = nd('Mul', [is_dom, a_flat], [[N]]) dom_sp = nd('Reshape', [is_dom_active, 'shape_hw'], [[1,1,H,W]]) # STEP 4: Bbox at active pixels mr_init = nd('Mul', ['row_s', dom_sp], [[1,1,H,W]]) max_row = prop_masked(mr_init, active) row_inv = nd('Sub', ['c_hm1', 'row_s'], [[1,1,H,W]]) mnr_init = nd('Mul', [row_inv, dom_sp], [[1,1,H,W]]) mnr_prop = prop_masked(mnr_init, active) min_row = nd('Sub', ['c_hm1', mnr_prop], [[1,1,H,W]]) mc_init = nd('Mul', ['col_s', dom_sp], [[1,1,H,W]]) max_col = prop_masked(mc_init, active) col_inv = nd('Sub', ['c_wm1', 'col_s'], [[1,1,H,W]]) mnc_init = nd('Mul', [col_inv, dom_sp], [[1,1,H,W]]) mnc_prop = prop_masked(mnc_init, active) min_col = nd('Sub', ['c_wm1', mnc_prop], [[1,1,H,W]]) # Flatten bbox min_row_f = nd('Reshape', [min_row, 'shape_n'], [[N]]) max_row_f = nd('Reshape', [max_row, 'shape_n'], [[N]]) min_col_f = nd('Reshape', [min_col, 'shape_n'], [[N]]) max_col_f = nd('Reshape', [max_col, 'shape_n'], [[N]]) h_comp_f = nd('Sub', [max_row_f, min_row_f], [[N]]) const('ones_n_c', np.ones(N, dtype=np.float32)) h_comp1_f = nd('Add', [h_comp_f, 'ones_n_c'], [[N]]) w_comp_f = nd('Sub', [max_col_f, min_col_f], [[N]]) w_comp1_f = nd('Add', [w_comp_f, 'ones_n'], [[N]]) local_r_f = nd('Sub', ['row_flat', min_row_f], [[N]]) local_c_f = nd('Sub', ['col_flat', min_col_f], [[N]]) # STEP 5: Per-color scatter not_dom_f = nd('Sub', ['ones_n', is_dom_active], [[N]]) painted_per_channel = [] for c in range(1, 10): ch_c = nd('Slice', ['input', f's_ch{c}', f'e_ch{c}', 'ax4'], [[1,1,H,W]]) ch_c_flat = nd('Reshape', [ch_c, 'shape_n'], [[N]]) minority_c = nd('Mul', [ch_c_flat, not_dom_f], [[N]]) mc_col = nd('Reshape', [minority_c, 'shape_n1'], [[N, 1]]) has_mc = nd('MatMul', [adj_m, mc_col], [[N, 1]]) has_mc_flat = nd('Reshape', [has_mc, 'shape_n'], [[N]]) has_mc_b = nd('Greater', [has_mc_flat, 'c_half'], [([N], TensorProto.BOOL)]) has_mc_f = nd('Cast', [has_mc_b], [[N]], to=1) row_mc = nd('Mul', ['row_flat', minority_c], [[N]]) row_mc_col = nd('Reshape', [row_mc, 'shape_n1'], [[N, 1]]) mc_row_sum = nd('MatMul', [adj_m, row_mc_col], [[N, 1]]) mc_row_f = nd('Reshape', [mc_row_sum, 'shape_n'], [[N]]) col_mc = nd('Mul', ['col_flat', minority_c], [[N]]) col_mc_col = nd('Reshape', [col_mc, 'shape_n1'], [[N, 1]]) mc_col_sum = nd('MatMul', [adj_m, col_mc_col], [[N, 1]]) mc_col_f = nd('Reshape', [mc_col_sum, 'shape_n'], [[N]]) mc_r_off = nd('Sub', [mc_row_f, min_row_f], [[N]]) gr_m_raw = nd('Div', [mc_r_off, h_comp1_f], [[N]]) gr_m = nd('Floor', [gr_m_raw], [[N]]) mc_c_off = nd('Sub', [mc_col_f, min_col_f], [[N]]) gc_m_raw = nd('Div', [mc_c_off, w_comp1_f], [[N]]) gc_m = nd('Floor', [gc_m_raw], [[N]]) const(f'two_n_{c}', np.full(N, 2.0, dtype=np.float32)) gr_d2 = nd('Div', [gr_m, f'two_n_{c}'], [[N]]) gr_fl = nd('Floor', [gr_d2], [[N]]) gr_2fl = nd('Mul', [gr_fl, f'two_n_{c}'], [[N]]) v_flip = nd('Sub', [gr_m, gr_2fl], [[N]]) gc_d2 = nd('Div', [gc_m, f'two_n_{c}'], [[N]]) gc_fl = nd('Floor', [gc_d2], [[N]]) gc_2fl = nd('Mul', [gc_fl, f'two_n_{c}'], [[N]]) h_flip = nd('Sub', [gc_m, gc_2fl], [[N]]) lr2 = nd('Mul', [local_r_f, f'two_n_{c}'], [[N]]) dr = nd('Sub', [h_comp_f, lr2], [[N]]) adj_r = nd('Mul', [v_flip, dr], [[N]]) dest_lr = nd('Add', [local_r_f, adj_r], [[N]]) lc2 = nd('Mul', [local_c_f, f'two_n_{c}'], [[N]]) dc = nd('Sub', [w_comp_f, lc2], [[N]]) adj_c = nd('Mul', [h_flip, dc], [[N]]) dest_lc = nd('Add', [local_c_f, adj_c], [[N]]) gr_h = nd('Mul', [gr_m, h_comp1_f], [[N]]) dest_r_off = nd('Add', [gr_h, dest_lr], [[N]]) dest_row = nd('Add', [min_row_f, dest_r_off], [[N]]) gc_w = nd('Mul', [gc_m, w_comp1_f], [[N]]) dest_c_off = nd('Add', [gc_w, dest_lc], [[N]]) dest_col = nd('Add', [min_col_f, dest_c_off], [[N]]) const(f'w_n_{c}', np.full(N, float(W), dtype=np.float32)) dest_rw = nd('Mul', [dest_row, f'w_n_{c}'], [[N]]) dest_idx_f = nd('Add', [dest_rw, dest_col], [[N]]) const(f'z_n_{c}', np.zeros(N, dtype=np.float32)) const(f'nm1_n_{c}', np.full(N, float(N-1), dtype=np.float32)) dest_idx_lo = nd('Max', [dest_idx_f, f'z_n_{c}'], [[N]]) dest_idx_clip = nd('Min', [dest_idx_lo, f'nm1_n_{c}'], [[N]]) dest_idx_i = nd('Cast', [dest_idx_clip], [([N], TensorProto.INT64)], to=7) scatter_vals = nd('Mul', [is_dom_active, has_mc_f], [[N]]) dest_idx_i_2d = nd('Reshape', [dest_idx_i, 'shape_n1'], [([N, 1], TensorProto.INT64)]) scatter_vals_2d = nd('Reshape', [scatter_vals, 'shape_n1'], [[N, 1]]) const(f'zeros_n1_{c}', np.zeros((N, 1), dtype=np.float32)) scattered = nd('ScatterElements', [f'zeros_n1_{c}', dest_idx_i_2d, scatter_vals_2d], [[N, 1]], axis=0, reduction='max') scattered_flat = nd('Reshape', [scattered, 'shape_n'], [[N]]) painted_per_channel.append(scattered_flat) # STEP 6: Build output const('zero_ch_flat', np.zeros(N, dtype=np.float32)) reshaped_channels = ['zero_ch_flat'] for pc in painted_per_channel: reshaped_channels.append(pc) ch_2d = [] for ch_name in reshaped_channels: r = nd('Reshape', [ch_name, 'shape_1n'], [[1, N]]) ch_2d.append(r) const('ax0', [0], 'i') painted_cn = nd('Concat', ch_2d, [[C, N]], axis=0) painted_1chw = nd('Reshape', [painted_cn, 'shape_1chw'], [[1, C, H, W]]) const('axes1b', [1], 'i') any_painted = nd('ReduceMax', [painted_1chw, 'axes1b'], [[1, 1, H, W]], keepdims=1) keep_mask = nd('Sub', ['c_one', any_painted], [[1, 1, H, W]]) kept = nd('Mul', ['input', keep_mask], [[1, C, H, W]]) nodes.append(helper.make_node('Add', [kept, painted_1chw], ['output'])) x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, C, H, W]) y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, C, H, W]) graph = helper.make_graph(nodes, 'task285_scatter', [x], [y], initializer=inits, value_info=vis) model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)]) return model def validate_with_official_utils(model_path, task_num, neurogolf_utils_path, task_data_dir): """ Validate and score model using the official Kaggle neurogolf_utils.py. This is the ONLY reliable way to verify correctness and get accurate scores. """ import sys import types # Mock IPython.display (not available outside Jupyter notebooks) 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 # Use non-interactive matplotlib backend import matplotlib matplotlib.use('Agg') # Import neurogolf_utils from the specified path utils_dir = os.path.dirname(os.path.abspath(neurogolf_utils_path)) if utils_dir not in sys.path: sys.path.insert(0, utils_dir) import neurogolf_utils # Override the data directory to local task data neurogolf_utils._NEUROGOLF_DIR = task_data_dir.rstrip('/') + '/' # Load examples using official loader examples = neurogolf_utils.load_examples(task_num) print(f" Loaded examples: train={len(examples['train'])}, " f"test={len(examples['test'])}, arc-gen={len(examples['arc-gen'])}") # Load model and run official verify_network network = onnx.load(model_path) print(f"\n{'='*60}") print(f" Official neurogolf_utils.verify_network() output:") print(f"{'='*60}\n") neurogolf_utils.verify_network(network, task_num, examples) print(f"\n{'='*60}") return True if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description="Build and validate optimized Task 285 model") parser.add_argument('--neurogolf-utils', default='../own-solver/neurogolf_utils.py', help='Path to official neurogolf_utils.py') parser.add_argument('--task-data-dir', default='../task-data', help='Directory containing taskNNN.json files (from neurogolf-2026.zip)') parser.add_argument('--output', default='/app/repo/medal-solvers/optimized/task285.onnx', help='Output path for the optimized model') parser.add_argument('--skip-official-validation', action='store_true', help='Skip official utils validation (NOT recommended)') args = parser.parse_args() # Auto-detect paths task_data_dir = args.task_data_dir if not os.path.exists(os.path.join(task_data_dir, 'task285.json')): for candidate in ['task-data', '../task-data', '.']: if os.path.exists(os.path.join(candidate, 'task285.json')): task_data_dir = candidate break neurogolf_path = args.neurogolf_utils if not os.path.exists(neurogolf_path): for candidate in ['../own-solver/neurogolf_utils.py', 'own-solver/neurogolf_utils.py', 'neurogolf_utils.py']: if os.path.exists(candidate): neurogolf_path = candidate break print("Building Task 285 optimized model (ScatterElements approach)...") model = build_task285() os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) onnx.save(model, args.output) fsize = os.path.getsize(args.output) print(f" Nodes: {len(model.graph.node)}") print(f" File size: {fsize:,} bytes (limit: {int(1.44*1024*1024):,})") print(f" Saved to: {args.output}") # Run official validation if not args.skip_official_validation: if not os.path.exists(neurogolf_path): print(f"\n WARNING: neurogolf_utils.py not found at {neurogolf_path}") print(f" Use --neurogolf-utils to specify path.") args.skip_official_validation = True elif not os.path.exists(os.path.join(task_data_dir, 'task285.json')): print(f"\n WARNING: task285.json not found in {task_data_dir}") print(f" Extract neurogolf-2026.zip first. Use --task-data-dir to specify.") args.skip_official_validation = True if not args.skip_official_validation: print(f"\n Using official neurogolf_utils.py from: {neurogolf_path}") print(f" Using task data from: {task_data_dir}") validate_with_official_utils(args.output, 285, neurogolf_path, task_data_dir) else: # Fallback: basic validation import onnxruntime import json print("\n Running basic validation (use --neurogolf-utils for official scoring)...") sess = onnxruntime.InferenceSession(args.output) task_json = os.path.join(task_data_dir, 'task285.json') if os.path.exists(task_json): with open(task_json) as f: td = json.load(f) right, wrong = 0, 0 for split in ['train', 'test', 'arc-gen']: for ex in td.get(split, []): 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 += 1 else: wrong += 1 print(f" Results: {right} pass, {wrong} fail") if wrong == 0: print(f" PASSED (use official utils for reliable score)") else: print(f" FAILED - do not submit!")