"""Build optimized ONNX model for Task 314. Rule (266/266): 8x8 grid divided into 3x3 arrangement of 2x2 blocks by 0-separators. For each non-1 color at a given relative position in blocks: if >=2 blocks in same row or column have that color → fill all blocks in that row/col. """ import numpy as np import onnx from onnx import helper, TensorProto, numpy_helper import onnxruntime as ort from onnxsim import simplify import json, os, math def build_task314(): 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 # Block gather indices: for (pr,pc), the 9 cells in flat-8x8 for pr in range(2): for pc in range(2): idx = [((bi*3+pr)*8 + (bj*3+pc)) for bi in range(3) for bj in range(3)] const(f'idx_{pr}{pc}', idx, 'i') # Output assembly: for each cell in 8x8, which position in combined[36]? # Combined = 4 grids of 9, indexed as (pr*2+pc)*9 + bi*3+bj cell_to_combined = np.zeros(64, dtype=np.int64) is_block = np.zeros(64, dtype=np.float32) for r in range(8): for c in range(8): if r in [2,5] or c in [2,5]: cell_to_combined[r*8+c] = 0 # dummy, will be masked is_block[r*8+c] = 0.0 else: if r < 2: bi, pr = 0, r elif r < 5: bi, pr = 1, r-3 else: bi, pr = 2, r-6 if c < 2: bj, pc = 0, c elif c < 5: bj, pc = 1, c-3 else: bj, pc = 2, c-6 flat_idx = (pr*2+pc)*9 + bi*3+bj cell_to_combined[r*8+c] = flat_idx is_block[r*8+c] = 1.0 const('cell_idx', cell_to_combined.tolist(), 'i') const('block_mask', is_block.reshape(1,1,64).tolist()) # [1,1,64] const('sl_s', [0,0,0,0], 'i') const('sl_e', [1,10,8,8], 'i') const('sl_ax', [0,1,2,3], 'i') # Slice for channels 2-9 (non-bg, non-base) const('ch_s', [0,2,0,0], 'i') const('ch_e', [1,10,8,8], 'i') const('shape_1_8_64', [1,8,64], 'i') const('shape_1_8_9', [1,8,9], 'i') const('shape_1_8_3_3', [1,8,3,3], 'i') const('shape_1_10_64', [1,10,64], 'i') const('shape_1_10_8_8', [1,10,8,8], 'i') const('shape_1_8_36', [1,8,36], 'i') const('shape_1_10_36', [1,10,36], 'i') const('c_1f', [1.0]) const('c_0f', [0.0]) const('axes2', [2], 'i') const('axes3', [3], 'i') const('pad_to_30', [0,0,0,0,0,0,22,22], 'i') const('pad_val', [0.0]) # Channel 1 one-hot for base color [1,10,1] ch1_oh = np.zeros((1,10,1), dtype=np.float32) ch1_oh[0,1,0] = 1.0 const('ch1_oh', ch1_oh) # Channel 0 one-hot for bg [1,10,1] ch0_oh = np.zeros((1,10,1), dtype=np.float32) ch0_oh[0,0,0] = 1.0 const('ch0_oh', ch0_oh) # Zeros for channels 0,1 in combined: [1,2,9] const('zeros_2_9', np.zeros((1,2,9), dtype=np.float32)) # Slice input to 8x8 inp8 = nd('Slice', ['input', 'sl_s', 'sl_e', 'sl_ax'], [[1,10,8,8]]) # Extract channels 2-9 [1,8,8,8] inp_nonbase = nd('Slice', [inp8, 'ch_s', 'ch_e', 'sl_ax'], [[1,8,8,8]]) inp_nb_flat = nd('Reshape', [inp_nonbase, 'shape_1_8_64'], [[1,8,64]]) # For each (pr,pc): gather, propagate filled_grids = [] for pr in range(2): for pc in range(2): gathered = nd('Gather', [inp_nb_flat, f'idx_{pr}{pc}'], [[1,8,9]], axis=2) g3x3 = nd('Reshape', [gathered, 'shape_1_8_3_3'], [[1,8,3,3]]) row_sum = nd('ReduceSum', [g3x3, 'axes3'], [[1,8,3,1]], keepdims=1) col_sum = nd('ReduceSum', [g3x3, 'axes2'], [[1,8,1,3]], keepdims=1) row_ge2 = nd('Cast', [nd('Less', ['c_1f', row_sum], [([1,8,3,1], TensorProto.BOOL)])], [[1,8,3,1]], to=1) col_ge2 = nd('Cast', [nd('Less', ['c_1f', col_sum], [([1,8,1,3], TensorProto.BOOL)])], [[1,8,1,3]], to=1) fill = nd('Max', [row_ge2, col_ge2], [[1,8,3,3]]) filled = nd('Max', [g3x3, fill], [[1,8,3,3]]) filled_flat = nd('Reshape', [filled, 'shape_1_8_9'], [[1,8,9]]) filled_grids.append(filled_flat) # Concat all 4 filled grids: [1,8,36] combined_8ch = nd('Concat', filled_grids, [[1,8,36]], axis=2) # Prepend zeros for channels 0,1: [1,2,36] + [1,8,36] → [1,10,36] const('zeros_2_36', np.zeros((1,2,36), dtype=np.float32)) combined_10ch = nd('Concat', ['zeros_2_36', combined_8ch], [[1,10,36]], axis=1) # Gather output: for each of 64 cells, pick from combined_10ch[36] block_vals = nd('Gather', [combined_10ch, 'cell_idx'], [[1,10,64]], axis=2) # For non-block cells, we need the original input values inp8_flat = nd('Reshape', [inp8, 'shape_1_10_64'], [[1,10,64]]) # block_mask: 1 for block cells, 0 for separators # output = block_mask * block_vals + (1-block_mask) * original # But for block cells: we need channels 0,1 too! # Block cells: ch 2-9 from propagation. ch 1 = 1 if no ch2-9 active. ch 0 = 0. # Separator cells: ch 0 = 1, others = 0. (original has this) # For block cells: compute ch1 = 1 - max(ch2-9) # Get ch2-9 sum per cell from block_vals[channels 2-9] ch2_9_vals = nd('Slice', [block_vals, 'ch_s', 'ch_e', 'sl_ax'], [[1,8,64]]) # wrong sl_ax # Actually need to slice on axis 1 (channels) const('ch_sl_s', [0,2,0], 'i') const('ch_sl_e', [1,10,64], 'i') const('ch_sl_ax', [0,1,2], 'i') ch2_9 = nd('Slice', [block_vals, 'ch_sl_s', 'ch_sl_e', 'ch_sl_ax'], [[1,8,64]]) # any_color = max over channels 2-9: [1,1,64] const('axes1', [1], 'i') any_color = nd('ReduceMax', [ch2_9, 'axes1'], [[1,1,64]], keepdims=1) # ch1_val = 1 - any_color (for block cells) [1,1,64] ch1_val = nd('Sub', ['c_1f', any_color], [[1,1,64]]) # Build full 10-channel block output: # ch0 = 0, ch1 = ch1_val, ch2-9 = ch2_9 const('zeros_1_1_64', np.zeros((1,1,64), dtype=np.float32)) block_full = nd('Concat', ['zeros_1_1_64', ch1_val, ch2_9], [[1,10,64]], axis=1) # Final: block_mask * block_full + (1-block_mask) * inp8_flat masked_block = nd('Mul', [block_full, 'block_mask'], [[1,10,64]]) const('inv_block_mask', (1.0 - is_block).reshape(1,1,64).tolist()) masked_orig = nd('Mul', [inp8_flat, 'inv_block_mask'], [[1,10,64]]) output_flat = nd('Add', [masked_block, masked_orig], [[1,10,64]]) # Reshape to 8x8 output_8x8 = nd('Reshape', [output_flat, 'shape_1_10_8_8'], [[1,10,8,8]]) # Pad to 30x30 final = nd('Pad', [output_8x8, 'pad_to_30', 'pad_val'], [[1,10,30,30]]) # Finish nodes.append(helper.make_node('Identity', [final], ['output'])) 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]) graph = helper.make_graph(nodes, 'task314', [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__': print("Building Task 314...") model = build_task314() onnx.save(model, '/app/task314_v2.onnx') # Simplify model_s, ok = simplify(model) if ok: del model_s.graph.value_info[:] model_s = onnx.shape_inference.infer_shapes(model_s, strict_mode=True) onnx.save(model_s, '/app/task314_opt.onnx') print(f"Simplified: {len(model_s.graph.node)} nodes, {os.path.getsize('/app/task314_opt.onnx'):,}B") else: print("Simplification failed, using original") onnx.save(model, '/app/task314_opt.onnx') # Validate with open('/app/task-data/task314.json') as f: data = json.load(f) sess = ort.InferenceSession('/app/task314_opt.onnx') all_ex = data['train'] + data['test'] + data.get('arc-gen', []) right, wrong = 0, 0 for ex in all_ex: inp_arr = 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_arr[0][v][r][c] = 1.0 result = sess.run(['output'], {'input': inp_arr})[0] out_pred = (result > 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_pred, exp): right += 1 else: wrong += 1 if wrong <= 3: for ch in range(10): d = np.where(out_pred[0,ch,:8,:8] != exp[0,ch,:8,:8]) if len(d[0]) > 0: print(f" Ch{ch}: {len(d[0])} diffs") print(f"Validation: {right}/{right+wrong}")