feat(optim): show custom operator domains

src/winml/modelkit/analyze/optim_output.py

@@ -127,11 +127,14 @@ def support_counts(self) -> dict[SupportLevel, int]:
     @staticmethod
     def _node_ref_dict(ref: NodeRef) -> dict[str, object]:
         """Return the stable JSON representation of one graph-delta node."""
-        return {
+        data: dict[str, object] = {
             "op_type": ref.op_type,
             "name": ref.name,
             "outputs": list(ref.outputs),
         }
+        if ref.domain and ref.domain != "ai.onnx":
+            data["domain"] = ref.domain
+        return data
 
     def to_dict(self) -> dict[str, object]:
         """Return actionable graph-delta and target-support evidence."""
@@ -239,7 +242,7 @@ def _check_one(
         support, reason = _lookup_support(ref, support_by_output)
         result.operators.append(
             ProducedOperatorSupport(
-                op_type=ref.op_type,
+                op_type=ref.qualified_op_type(),
                 label=ref.label(),
                 change=change,
                 support=support,

src/winml/modelkit/optim/analysis.py

@@ -56,11 +56,19 @@ class NodeRef:
         op_type: The node's operator type (e.g. ``"MatMul"``).
         name: The node's name (may be empty — ONNX names are optional).
         outputs: The node's output tensor names.
+        domain: The node's operator domain. Empty means the default ONNX domain.
     """
 
     op_type: str
     name: str
     outputs: tuple[str, ...]
+    domain: str = ""
+
+    def qualified_op_type(self) -> str:
+        """Return the operator type, qualified when it uses a custom domain."""
+        if self.domain and self.domain != "ai.onnx":
+            return f"{self.domain}::{self.op_type}"
+        return self.op_type
 
     def label(self) -> str:
         """Return a human-readable identifier for this node.
@@ -69,7 +77,7 @@ def label(self) -> str:
         since output names are unique within a graph.
         """
         ident = self.name or (self.outputs[0] if self.outputs else "?")
-        return f"{self.op_type} '{ident}'"
+        return f"{self.qualified_op_type()} '{ident}'"
 
 
 @dataclass
@@ -136,7 +144,7 @@ def op_histogram(self, kind: str) -> list[tuple[str, int]]:
             "added": self.added_nodes,
             "modified": self.modified_nodes,
         }[kind]
-        return Counter(n.op_type for n in nodes).most_common()
+        return Counter(n.qualified_op_type() for n in nodes).most_common()
 
 
 # =============================================================================
@@ -163,7 +171,7 @@ def _node_identity(node: NodeProto) -> tuple[Any, ...]:
     """
     if len(node.output) > 0:
         return tuple(node.output)
-    return ("\0no-output", node.op_type, node.name, tuple(node.input))
+    return ("\0no-output", node.domain, node.op_type, node.name, tuple(node.input))
 
 
 def _collect_nodes(
@@ -186,7 +194,7 @@ def _collect_nodes(
             key = (cur_scope, _node_identity(node))
             table[key] = (
                 node.SerializeToString(),
-                NodeRef(node.op_type, node.name, tuple(node.output)),
+                NodeRef(node.op_type, node.name, tuple(node.output), node.domain),
             )
             for attr in node.attribute:
                 if attr.type == AttributeProto.GRAPH:

tests/unit/analyze/test_optim_output.py

@@ -174,11 +174,16 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None:
             category="rewrite",
             description="Replace static Split with Slice.",
             pipe_name="algebraic",
-            removed_nodes=[NodeRef("Split", "split", ("a", "b"))],
-            added_nodes=[NodeRef("Slice", "slice_0", ("a",))],
+            removed_nodes=[NodeRef("Split", "split", ("a", "b"), "ai.onnx")],
+            added_nodes=[NodeRef("Slice", "slice_0", ("a",), "com.microsoft")],
             modified_initializers=["starts"],
             operators=[
-                ProducedOperatorSupport("Slice", "Slice 'slice_0'", "added", SupportLevel.SUPPORTED)
+                ProducedOperatorSupport(
+                    "com.microsoft::Slice",
+                    "com.microsoft::Slice 'slice_0'",
+                    "added",
+                    SupportLevel.SUPPORTED,
+                )
             ],
         )
 
@@ -192,16 +197,23 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None:
             "support_counts": {"supported": 1},
             "graph_delta": {
                 "removed_nodes": [{"op_type": "Split", "name": "split", "outputs": ["a", "b"]}],
-                "added_nodes": [{"op_type": "Slice", "name": "slice_0", "outputs": ["a"]}],
+                "added_nodes": [
+                    {
+                        "op_type": "Slice",
+                        "name": "slice_0",
+                        "outputs": ["a"],
+                        "domain": "com.microsoft",
+                    }
+                ],
                 "modified_nodes": [],
                 "removed_initializers": [],
                 "added_initializers": [],
                 "modified_initializers": ["starts"],
             },
             "operators": [
                 {
-                    "op_type": "Slice",
-                    "label": "Slice 'slice_0'",
+                    "op_type": "com.microsoft::Slice",
+                    "label": "com.microsoft::Slice 'slice_0'",
                     "change": "added",
                     "support": "supported",
                 }

tests/unit/commands/test_optimize_cli.py

@@ -276,7 +276,9 @@ def test_device_target_forwarded_to_optimizer(
 _ANALYZE_MODEL = "winml.modelkit.optim.analyze_model"
 
 
-def _make_finding(name: str = "clamp-constant-values") -> MagicMock:
+def _make_finding(
+    name: str = "clamp-constant-values", node_domain: str = ""
+) -> MagicMock:
     """Build a stand-in CapabilityFinding for renderer tests."""
     from winml.modelkit.optim import CapabilityFinding, NodeRef
 
@@ -288,7 +290,7 @@ def _make_finding(name: str = "clamp-constant-values") -> MagicMock:
         description="clamp things",
         pipe_name="surgery",
         modified_initializers=["BIG"],
-        removed_nodes=[NodeRef("MatMul", "mm", ("mm",))],
+        removed_nodes=[NodeRef("MatMul", "mm", ("mm",), node_domain)],
     )
 
 
@@ -335,6 +337,24 @@ def test_check_optim_lists_applicable_flag(self, runner: CliRunner, tmp_path: Pa
         assert "--enable-matmul-add-fusion" in result.output
         assert "1 applicable optimization" in result.output
 
+    def test_check_optim_shows_custom_operator_domain(
+        self, runner: CliRunner, tmp_path: Path
+    ) -> None:
+        model_file = tmp_path / "model.onnx"
+        model_file.touch()
+
+        with (
+            patch(_LOAD_ONNX, return_value=_make_mock_model()),
+            patch(
+                _ANALYZE_MODEL,
+                return_value=[_make_finding(node_domain="com.microsoft")],
+            ),
+        ):
+            result = runner.invoke(optimize, ["-m", str(model_file), "--check-optim"])
+
+        assert result.exit_code == 0, result.output
+        assert "com.microsoft::MatMul" in result.output
+
     def test_check_optim_no_findings_message(self, runner: CliRunner, tmp_path: Path) -> None:
         model_file = tmp_path / "model.onnx"
         model_file.touch()

tests/unit/optim/test_analysis.py

@@ -137,6 +137,29 @@ def test_identical_graphs_have_no_diff(self) -> None:
         removed, added, modified = _diff_nodes(table_a, table_b)
         assert not removed and not added and not modified
 
+    def test_collected_node_preserves_custom_domain(self) -> None:
+        graph = helper.make_graph(
+            [
+                helper.make_node(
+                    "Gelu",
+                    ["x"],
+                    ["y"],
+                    name="gelu",
+                    domain="com.microsoft",
+                )
+            ],
+            "custom_domain",
+            [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])],
+            [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])],
+        )
+        table: dict = {}
+
+        _collect_nodes(graph, (), table)
+
+        ref = next(iter(table.values()))[1]
+        assert ref.domain == "com.microsoft"
+        assert ref.qualified_op_type() == "com.microsoft::Gelu"
+
     def test_subgraph_nodes_are_collected(self) -> None:
         """Nodes inside a control-flow subgraph are included in the table."""
         then_graph = helper.make_graph(
@@ -505,6 +528,32 @@ def test_node_ref_label_uses_name_then_output(self) -> None:
         assert NodeRef("MatMul", "mm", ("mm_out",)).label() == "MatMul 'mm'"
         assert NodeRef("Add", "", ("y",)).label() == "Add 'y'"
 
+    def test_node_ref_label_qualifies_custom_domains(self) -> None:
+        assert (
+            NodeRef("Gelu", "gelu", ("y",), "com.microsoft").label()
+            == "com.microsoft::Gelu 'gelu'"
+        )
+        assert NodeRef("Add", "add", ("y",), "ai.onnx").label() == "Add 'add'"
+
+    def test_op_histogram_distinguishes_custom_domains(self) -> None:
+        finding = CapabilityFinding(
+            name="x",
+            python_name="x",
+            enable_flag="--enable-x",
+            category="misc",
+            description="",
+            pipe_name="p",
+            added_nodes=[
+                NodeRef("Gelu", "standard", ("a",)),
+                NodeRef("Gelu", "contrib", ("b",), "com.microsoft"),
+            ],
+        )
+
+        assert finding.op_histogram("added") == [
+            ("Gelu", 1),
+            ("com.microsoft::Gelu", 1),
+        ]
+
 
 # =============================================================================
 # END-TO-END ANALYSIS