Create build_task363_onnx.py

medal-solvers/build_tasks/build_task363_onnx.py

@@ -0,0 +1,194 @@
+
+"""
+Compact ONNX model for ARC task 363.
+
+Rule (verified on all 265 train/test/arc-gen examples; grids are all 10x10):
+  The input has a "template" drawn in colour 2 (within a box of colour-5 background) and
+  the rest of the grid is a colour-0 / colour-5 texture.  Stamp the template (its 2-cells)
+  at every placement where the 2-cells land on colour-0 cells, subject to filtering:
+    valid       : all template-2 cells align with colour-0 cells.
+    has_overlap : the placement's painted cells overlap another placement's.
+    touch_edge  : the placement reaches the grid border.
+    dense       : (fg dilated by the template-shape) count > 16 at that offset.
+    top2        : the placement lies entirely within the top two rows.
+    keep = valid AND (no_overlap OR touch_edge)
+                 AND NOT (no_overlap AND not touch_edge AND dense AND top2)
+  Painted cells become colour 2; everything else unchanged.
+
+Memory strategy (few intermediates, small dtypes, work in 10x10 / 19x19 offset space):
+  All channel reads come from a single colour-value map produced by one 1x1 Conv.
+  Placement search / overlap / edge / density / top2 are Conv / ConvTranspose in the
+  19x19 offset space (float16).  Output is composed with one Equal at the end.
+"""
+import numpy as np
+import onnx
+from onnx import helper, TensorProto, numpy_helper
+
+C = 10
+G = 10          # actual grid size (all examples are 10x10)
+F = TensorProto.FLOAT16
+I64 = TensorProto.INT64
+BOOL = TensorProto.BOOL
+
+
+def f16(name, arr):
+    return numpy_helper.from_array(np.asarray(arr, dtype=np.float16), name=name)
+
+
+def i64(name, arr):
+    return numpy_helper.from_array(np.asarray(arr, dtype=np.int64), name=name)
+
+
+def build_model():
+    n = []
+    I = []
+
+    # ---- constants ----
+    # read channel 2 and channel 5 (the only non-zero colours) as thin 10x10 slices
+    I.append(i64("sp_ax", [2, 3]))
+    I.append(i64("ch_ax", [1, 2, 3]))
+    I.append(i64("c2s", [2, 0, 0])); I.append(i64("c2e", [3, G, G]))
+    I.append(i64("c5s", [5, 0, 0])); I.append(i64("c5e", [6, G, G]))
+    I.append(f16("five16", [5.0]))
+    I.append(f16("v2", [2.0]))
+    I.append(f16("half", [0.5]))
+    I.append(f16("d16", [16.0]))
+    # colour-value 1x1 conv: cval = 1*ch0 + 2*ch2 + 5*ch5 ... actually we need: is2, is0, isFG.
+    # Use weights so cval encodes: ch0->1, ch2->2, ch5->5 (others 0). Then is2=(cval==2),
+    # is0=(cval==1)? no. Simpler: produce 3 derived maps with one Conv of 3 output channels:
+    #   out0 = ch2  (template), out1 = ch0 (zero), out2 = sum(ch1..ch9) (foreground).
+    # 1x1 Conv reads only channels 2..5 (cval = 2*ch2 + 5*ch5; bg/zero => cval==0).
+    I.append(f16("one16", [1.0]))
+
+    # k3 ones for dilation
+    I.append(f16("k3", np.ones((1, 1, 3, 3), np.float16)))
+    # edge ring (10x10) and top-2-rows (10x10)
+    edge = np.zeros((1, 1, G, G), np.float16)
+    edge[0, 0, 0, :] = 1; edge[0, 0, -1, :] = 1; edge[0, 0, :, 0] = 1; edge[0, 0, :, -1] = 1
+    I.append(f16("edge10", edge))
+    top2 = np.zeros((1, 1, G, G), np.float16); top2[0, 0, 0:2, :] = 1
+    I.append(f16("top2_10", top2))
+    # pad 10 -> 28 (9 each side) for "full" convs with a 10x10 kernel
+    # crop 28 -> 10 (centre) for the painted result
+    # pad 10 -> 30 to restore full output grid
+    # output one-hot: uint8 colour index vs channel range
+    U8 = TensorProto.UINT8
+    I.append(numpy_helper.from_array(np.arange(C, dtype=np.uint8).reshape(1, C, 1, 1), "chan_vals_u8"))
+    I.append(numpy_helper.from_array(np.array(2, np.uint8).reshape(1, 1, 1, 1), "u2"))
+    I.append(numpy_helper.from_array(np.array(255, np.uint8).reshape(1, 1, 1, 1), "u255"))
+    # pad 10 -> 30 with sentinel 255 for the colour-index map
+    I.append(i64("pad_to30_idx", [0, 0, 0, 0, 0, 0, 20, 20]))
+
+    # ---- read only channel-2 (template) and channel-5 (texture) as thin 10x10 slices ----
+    n.append(helper.make_node("Slice", ["input", "c2s", "c2e", "ch_ax"], ["T_f32"]))   # [1,1,10,10] f32 ch2
+    n.append(helper.make_node("Slice", ["input", "c5s", "c5e", "ch_ax"], ["F5_f32"]))  # [1,1,10,10] f32 ch5
+    n.append(helper.make_node("Cast", ["T_f32"], ["T"], to=F))                      # template (colour 2)
+    n.append(helper.make_node("Cast", ["F5_f32"], ["F5"], to=F))                    # texture (colour 5)
+    n.append(helper.make_node("Add", ["T", "F5"], ["FG"]))                          # foreground = ch2|ch5 (mutually excl.)
+    n.append(helper.make_node("Sub", ["one16", "FG"], ["Z"]))                       # zero = 1 - fg
+    # colour value for output: 2*T + 5*F5
+    n.append(helper.make_node("Mul", ["F5", "five16"], ["c5v"]))
+    n.append(helper.make_node("Mul", ["T", "v2"], ["c2v"]))
+    n.append(helper.make_node("Add", ["c2v", "c5v"], ["cval"]))                     # 0 / 2 / 5
+    n.append(helper.make_node("ReduceSum", ["T", "sp_ax"], ["nt"], keepdims=1))     # template size
+
+    # ---- valid placements: Conv(zero, T, pad9) == nt  (19x19 offset space) ----
+    n.append(helper.make_node("Conv", ["Z", "T"], ["match"], pads=[9, 9, 9, 9]))    # [1,1,19,19]
+    n.append(helper.make_node("Equal", ["match", "nt"], ["valid_b"]))
+    n.append(helper.make_node("Cast", ["valid_b"], ["valid"], to=F))
+
+    # ---- overlap: ConvTranspose(valid, T, pads=9) -> cover[10x10]; Conv(cover, T, pad9) -> ov[19]; >nt ----
+    n.append(helper.make_node("ConvTranspose", ["valid", "T"], ["cover"], pads=[9, 9, 9, 9]))  # [1,1,10,10]
+    n.append(helper.make_node("Conv", ["cover", "T"], ["ov"], pads=[9, 9, 9, 9]))   # [1,1,19,19]
+    n.append(helper.make_node("Greater", ["ov", "nt"], ["has_ov"]))
+    n.append(helper.make_node("Not", ["has_ov"], ["no_ov"]))
+
+    # ---- touch edge: Conv(edge, T, pad9) > 0.5 ----
+    n.append(helper.make_node("Conv", ["edge10", "T"], ["ec"], pads=[9, 9, 9, 9]))
+    n.append(helper.make_node("Greater", ["ec", "half"], ["touch"]))
+    n.append(helper.make_node("Not", ["touch"], ["not_touch"]))
+    n.append(helper.make_node("Or", ["no_ov", "touch"], ["keep_ov"]))
+
+    # ---- dense: template_dilated = Conv(T, k3, pad1); Conv(fg_pad, Td) > 16 ----
+    n.append(helper.make_node("Conv", ["T", "k3"], ["Td"], pads=[1, 1, 1, 1]))      # [1,1,10,10]
+    n.append(helper.make_node("Conv", ["FG", "Td"], ["nz8"], pads=[9, 9, 9, 9]))    # [1,1,19,19]
+    n.append(helper.make_node("Greater", ["nz8", "d16"], ["dense"]))
+
+    # ---- top2: Conv(top2, T, pad9) == nt ----
+    n.append(helper.make_node("Conv", ["top2_10", "T"], ["t2c"], pads=[9, 9, 9, 9]))
+    n.append(helper.make_node("Equal", ["t2c", "nt"], ["all_top2"]))
+
+    # ---- top_false = (no_ov & not_touch) & (dense & all_top2) ; kept = valid & keep_ov & ~top_false
+    n.append(helper.make_node("And", ["dense", "all_top2"], ["dat"]))
+    n.append(helper.make_node("And", ["no_ov", "not_touch"], ["nont"]))
+    n.append(helper.make_node("And", ["dat", "nont"], ["top_false"]))
+    n.append(helper.make_node("Not", ["top_false"], ["not_tf"]))
+    n.append(helper.make_node("And", ["valid_b", "keep_ov"], ["kept_a"]))
+    n.append(helper.make_node("And", ["kept_a", "not_tf"], ["kept_b"]))
+    n.append(helper.make_node("Cast", ["kept_b"], ["kept"], to=F))
+
+    # ---- paint: ConvTranspose(kept, T, pads=9) -> 10x10 directly (no 28x28, no crop) ----
+    n.append(helper.make_node("ConvTranspose", ["kept", "T"], ["paint10"], pads=[9, 9, 9, 9]))
+    n.append(helper.make_node("Greater", ["paint10", "half"], ["paint_b"]))  # [1,1,10,10] bool
+
+    # ---- compose output via a uint8 colour-index map (one small tensor + one Equal) ----
+    #   colour_idx (10x10) = 2 where painted, else cval (0/2/5).  Pad to 30x30 with 255
+    #   (out-of-grid) so the final Equal yields all-False there.  output one-hot = Equal.
+    n.append(helper.make_node("Cast", ["cval"], ["cval_u8"], to=U8))                 # [1,1,10,10] u8
+    n.append(helper.make_node("Where", ["paint_b", "u2", "cval_u8"], ["cidx10"]))    # [1,1,10,10] u8
+    n.append(helper.make_node("Pad", ["cidx10", "pad_to30_idx", "u255"], ["cidx30"]))  # [1,1,30,30] u8
+    n.append(helper.make_node("Equal", ["chan_vals_u8", "cidx30"], ["output"]))      # [1,C,30,30] bool
+
+    inp = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, C, 30, 30])
+    out = helper.make_tensor_value_info("output", BOOL, [1, C, 30, 30])
+    g = helper.make_graph(n, "task363", [inp], [out], I)
+    m = helper.make_model(g, opset_imports=[helper.make_operatorsetid("", 18)], ir_version=10)
+    onnx.checker.check_model(m)
+    return m
+
+
+def encode_grid(grid):
+    g = np.array(grid, dtype=np.int64)
+    h, w = g.shape
+    oh = np.zeros((1, C, 30, 30), dtype=np.float32)
+    for r in range(h):
+        for c in range(w):
+            oh[0, int(g[r, c]), r, c] = 1.0
+    return oh
+
+
+def verify(onnx_path, task_json):
+    import json, onnxruntime as ort
+    data = json.load(open(task_json))
+    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+    total = passed = 0
+    fails = []
+    for split in ("train", "test", "arc-gen"):
+        for i, pair in enumerate(data.get(split, [])):
+            inp, tgt = pair["input"], pair["output"]
+            if max(len(inp), max((len(r) for r in inp), default=0)) > 30:
+                continue
+            y = sess.run(["output"], {"input": encode_grid(inp)})[0]
+            if np.array_equal((y > 0.0).astype(np.float32), encode_grid(tgt)):
+                passed += 1
+            else:
+                fails.append((split, i))
+            total += 1
+    print(f"correctness: {passed}/{total} pass" + (f"  FAILS={fails[:10]}" if fails else ""))
+    return passed == total
+
+
+if __name__ == "__main__":
+    import argparse, os
+    ap = argparse.ArgumentParser(description="Build and verify Task 363 ONNX model.")
+    ap.add_argument("--input-json", default="/home/user/uploads/task363.json")
+    ap.add_argument("--output-onnx", default="task363.onnx")
+    ap.add_argument("--no-verify", action="store_true")
+    args = ap.parse_args()
+    m = build_model()
+    onnx.save(m, args.output_onnx)
+    print(f"saved {args.output_onnx}: {os.path.getsize(args.output_onnx)} bytes, "
+          f"{len(m.graph.node)} nodes")
+    if not args.no_verify and os.path.exists(args.input_json):
+        verify(args.output_onnx, args.input_json)
+