"""Build optimized ONNX for Task 398 — fixed version (no dynamic N). Rule (268/268 verified): - Input is 1x5 with 1-5 non-zero colored values - Output is NxN where N = count_nonzero * 5 - Diagonal waterfall: out[r,c] = input[0, c+r+1-N] if index valid and non-zero ONNX approach (state-safe): - For each possible N (5,10,15,20,25), precompute the output using FIXED offsets - Count non-zero from input - Select the correct branch based on nz_count - This avoids dynamic intermediate values that cause state bleeding """ 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_task398(): 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_4_5', [4.5]) const('c_neg_half', [-0.5]) const('axes1', [1], 'i') const('axes23', [2, 3], 'i') const('shape_1_10_1_1', [1, 10, 1, 1], 'i') const('shape_1_1_1_1', [1, 1, 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') # === STEP 1: Count non-zero === ch0 = nd('Gather', ['input', 'idx_0'], [[1, 1, 30, 30]], axis=1) 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 = nd('Mul', [active, nd('Sub', ['c_one', ch0], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]) nz_count = nd('ReduceSum', [fg, 'axes23'], [[1, 1, 1, 1]], keepdims=1) # === STEP 2: For each N, build the output === const('zeros_1_10_30_30', np.zeros((1, 10, 30, 30), dtype=np.float32)) def build_branch_for_N(N_val): """Build output for a specific N value using CONSTANT offsets.""" # src_col = col + row + 1 - N (all constant operations on grids) # Precompute as constant src_col_arr = col_grid + row_grid + (1.0 - N_val) # [1,1,30,30] constant offset const(f'src_col_N{N_val}', src_col_arr.astype(np.float32)) src_col_f = f'src_col_N{N_val}' # Valid: src in [0,4] src_valid_lo = nd('Greater', [src_col_f, 'c_neg_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) src_valid_hi = nd('Less', [src_col_f, 'c_4_5'], [([1, 1, 30, 30], TensorProto.BOOL)]) src_valid = nd('And', [src_valid_lo, src_valid_hi], [([1, 1, 30, 30], TensorProto.BOOL)]) src_valid_f = nd('Cast', [src_valid], [[1, 1, 30, 30]], to=1) # Within NxN grid const(f'N_minus_half_{N_val}', [N_val - 0.5]) in_grid_r = nd('Less', ['row_grid', f'N_minus_half_{N_val}'], [([1, 1, 30, 30], TensorProto.BOOL)]) in_grid_c = nd('Less', ['col_grid', f'N_minus_half_{N_val}'], [([1, 1, 30, 30], TensorProto.BOOL)]) in_grid = nd('And', [in_grid_r, in_grid_c], [([1, 1, 30, 30], TensorProto.BOOL)]) in_grid_f = nd('Cast', [in_grid], [[1, 1, 30, 30]], to=1) valid = nd('Mul', [src_valid_f, in_grid_f], [[1, 1, 30, 30]]) # For each source position, check and gather color result = 'zeros_1_10_30_30' for pos in range(5): const(f'pos_{pos}_N{N_val}', [float(pos)]) src_eq = nd('Less', [nd('Abs', [nd('Sub', [src_col_f, f'pos_{pos}_N{N_val}'], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]), 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) src_eq_f = nd('Cast', [src_eq], [[1, 1, 30, 30]], to=1) pos_mask = nd('Mul', [src_eq_f, valid], [[1, 1, 30, 30]]) # Get input color at (0, pos) const(f'sl_s_{pos}_N{N_val}', [0, 0, 0, pos], 'i') const(f'sl_e_{pos}_N{N_val}', [1, 10, 1, pos+1], 'i') const(f'sl_a_{pos}_N{N_val}', [0, 1, 2, 3], 'i') input_at_pos = nd('Slice', ['input', f'sl_s_{pos}_N{N_val}', f'sl_e_{pos}_N{N_val}', f'sl_a_{pos}_N{N_val}'], [[1, 10, 1, 1]]) colored = nd('Mul', [pos_mask, input_at_pos], [[1, 10, 30, 30]]) result = nd('Add', [result, colored], [[1, 10, 30, 30]]) # Add background result_sum = nd('ReduceSum', [result, 'axes1'], [[1, 1, 30, 30]], keepdims=1) has_result = nd('Greater', [result_sum, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)]) has_result_f = nd('Cast', [has_result], [[1, 1, 30, 30]], to=1) bg_mask = nd('Sub', [in_grid_f, has_result_f], [[1, 1, 30, 30]]) bg_mask_clip = nd('Clip', [bg_mask, 'c_zero', 'c_one'], [[1, 1, 30, 30]]) 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]]) bg_col = nd('Mul', [bg_mask_clip, ch0_oh_4d], [[1, 10, 30, 30]]) return nd('Add', [result, bg_col], [[1, 10, 30, 30]]) # Build all 5 branches branches = [] for k in range(1, 6): N_val = k * 5 branch = build_branch_for_N(N_val) branches.append(branch) # === STEP 3: Select based on nz_count === final = 'zeros_1_10_30_30' for k in range(1, 6): const(f'k_val_{k}', [float(k)]) # gate = (nz_count == k) diff_k = nd('Abs', [nd('Sub', [nz_count, f'k_val_{k}'], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]) eq_k = nd('Less', [diff_k, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)]) gate_k = nd('Cast', [eq_k], [[1, 1, 1, 1]], to=1) gated = nd('Mul', [branches[k-1], gate_k], [[1, 10, 30, 30]]) final = nd('Add', [final, gated], [[1, 10, 30, 30]]) return b.finish('task398', last_tensor=final) if __name__ == '__main__': import os, json, math, types import numpy as np import onnx import onnxruntime as ort print("Building Task 398 ONNX model (fixed)...") model = build_task398() os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True) output_path = '/app/repo/medal-solvers/optimized/task398.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/task398.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: print(f" FAIL {i}") 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 = '398' sess_off = ort.InferenceSession(sanitized.SerializeToString(), options) examples = neurogolf_utils.load_examples(398) 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})")