"""Build ONNX model for Task 277 — Smallest CC recoloring. V2 compact. Key facts: Always exactly 3 CCs, gap≥2, no ties. Approach: 1. MaxPool label propagation → 3 unique labels 2. Find 3 labels: max, second max, third max 3. Count cells per label 4. Min count → color 2, others → color 1 """ import numpy as np import onnx from onnx import helper, numpy_helper, TensorProto import math import os def build_task277(): 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_one', [1.0]) const('c_zero', [0.0]) const('c_half', [0.5]) const('c_big', [200.0]) 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('sl_ch8_start', [0, 8, 0, 0], 'i') const('sl_ch8_end', [1, 9, 10, 10], 'i') # Label grid: unique per cell (1-100) const('label_init', np.arange(1, 101, dtype=np.float32).reshape(1, 1, 10, 10)) const('axes23', [2, 3], 'i') const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i') const('pad_val_zero', [0.0]) const('zeros_7ch', np.zeros((1, 7, 10, 10), dtype=np.float32)) # === STEP 1: Get mask === inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]]) mask8 = nd('Slice', [inp10, 'sl_ch8_start', 'sl_ch8_end', 'sl_axes'], [[1, 1, 10, 10]]) # === STEP 2: Label propagation === labels = nd('Mul', [mask8, 'label_init'], [[1, 1, 10, 10]]) for _ in range(10): pooled = nd('MaxPool', [labels], [[1, 1, 10, 10]], kernel_shape=[3, 3], pads=[1, 1, 1, 1]) labels = nd('Mul', [pooled, mask8], [[1, 1, 10, 10]]) # === STEP 3: Find 3 unique labels === # L1 = max of all labels L1 = nd('ReduceMax', [labels, 'axes23'], [[1, 1, 1, 1]], keepdims=1) # Mask out L1: set cells with label==L1 to 0 diff_L1 = nd('Abs', [nd('Sub', [labels, L1], [[1, 1, 10, 10]])], [[1, 1, 10, 10]]) is_L1_b = nd('Less', [diff_L1, 'c_half'], [([1, 1, 10, 10], TensorProto.BOOL)]) is_L1 = nd('Cast', [is_L1_b], [[1, 1, 10, 10]], to=1) not_L1 = nd('Sub', ['c_one', is_L1], [[1, 1, 10, 10]]) labels_no_L1 = nd('Mul', [labels, not_L1], [[1, 1, 10, 10]]) # L2 = max of remaining labels L2 = nd('ReduceMax', [labels_no_L1, 'axes23'], [[1, 1, 1, 1]], keepdims=1) # Mask out L2 diff_L2 = nd('Abs', [nd('Sub', [labels, L2], [[1, 1, 10, 10]])], [[1, 1, 10, 10]]) is_L2_b = nd('Less', [diff_L2, 'c_half'], [([1, 1, 10, 10], TensorProto.BOOL)]) is_L2 = nd('Cast', [is_L2_b], [[1, 1, 10, 10]], to=1) not_L2 = nd('Sub', ['c_one', is_L2], [[1, 1, 10, 10]]) labels_no_L1_L2 = nd('Mul', [labels_no_L1, not_L2], [[1, 1, 10, 10]]) # L3 mask is just: mask8 - is_L1 - is_L2 is_L3 = nd('Sub', [mask8, nd('Add', [is_L1, is_L2], [[1, 1, 10, 10]])], [[1, 1, 10, 10]]) # === STEP 4: Count cells per CC === count_L1 = nd('ReduceSum', [is_L1, 'axes23'], [[1, 1, 1, 1]], keepdims=1) count_L2 = nd('ReduceSum', [is_L2, 'axes23'], [[1, 1, 1, 1]], keepdims=1) count_L3 = nd('ReduceSum', [is_L3, 'axes23'], [[1, 1, 1, 1]], keepdims=1) # === STEP 5: Find minimum count === min_12 = nd('Min', [count_L1, count_L2], [[1, 1, 1, 1]]) min_count = nd('Min', [min_12, count_L3], [[1, 1, 1, 1]]) # Which CC has this count? (no ties guaranteed) diff_c1 = nd('Abs', [nd('Sub', [count_L1, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]) smallest_is_L1_b = nd('Less', [diff_c1, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)]) smallest_is_L1 = nd('Cast', [smallest_is_L1_b], [[1, 1, 1, 1]], to=1) diff_c2 = nd('Abs', [nd('Sub', [count_L2, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]) smallest_is_L2_b = nd('Less', [diff_c2, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)]) smallest_is_L2 = nd('Cast', [smallest_is_L2_b], [[1, 1, 1, 1]], to=1) diff_c3 = nd('Abs', [nd('Sub', [count_L3, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]]) smallest_is_L3_b = nd('Less', [diff_c3, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)]) smallest_is_L3 = nd('Cast', [smallest_is_L3_b], [[1, 1, 1, 1]], to=1) # === STEP 6: Build smallest mask and other mask === sm1 = nd('Mul', [is_L1, smallest_is_L1], [[1, 1, 10, 10]]) sm2 = nd('Mul', [is_L2, smallest_is_L2], [[1, 1, 10, 10]]) sm3 = nd('Mul', [is_L3, smallest_is_L3], [[1, 1, 10, 10]]) smallest_mask = nd('Add', [sm1, nd('Add', [sm2, sm3], [[1, 1, 10, 10]])], [[1, 1, 10, 10]]) other_mask = nd('Sub', [mask8, smallest_mask], [[1, 1, 10, 10]]) # === STEP 7: Build output === bg_mask = nd('Sub', ['c_one', mask8], [[1, 1, 10, 10]]) out_10 = nd('Concat', [bg_mask, other_mask, smallest_mask, 'zeros_7ch'], [[1, 10, 10, 10]], axis=1) # === STEP 8: Pad — output directly === nodes.append(helper.make_node('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], ['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, 'task277', [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 print("Building Task 277 ONNX model v2...") model = build_task277() output_path = '/app/task277.onnx' del model.graph.value_info[:] model = onnx.shape_inference.infer_shapes(model, strict_mode=True) onnx.save(model, output_path) print(f" Nodes: {len(model.graph.node)}") print(f" File size: {os.path.getsize(output_path):,} bytes") sess = ort.InferenceSession(output_path) with open('/app/task-data/task277.json') as f: data = json.load(f) 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 <= 5: print(f" FAIL {i}: {len(np.where(out != exp)[0])} diffs") print(f" Results: {right_count}/{right_count+wrong_count} pass ({wrong_count} fail)") 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])) * np.dtype(onnx.helper.tensor_dtype_to_np_dtype(vi.type.tensor_type.elem_type)).itemsize 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: {mem_est:,}, Score: {score_est:.3f}") print(f" Base: 13.30. Gain: {score_est - 13.30:+.3f}")