improves dot rendering

Author: xadupre

_unittests/ut_helpers/test_dot_helper.py

@@ -62,6 +62,63 @@ def test_custom_doc_kernels_layer_normalization(self):
         self.maxDiff = None
         self.assertEqual(expected.strip("\n "), dot.strip("\n "))
 
+    def test_custom_doc_kernels_layer_normalization_constant(self):
+        TFLOAT16 = onnx.TensorProto.FLOAT16
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node(
+                        "LayerNormalization",
+                        ["X", "W", "B"],
+                        ["ln"],
+                        axis=-1,
+                        epsilon=9.999999974752427e-7,
+                    ),
+                    oh.make_node("Constant", [], ["cst"], value_float=[1]),
+                    oh.make_node("Cast", ["cst"], ["cst16"], to=onnx.TensorProto.FLOAT16),
+                    oh.make_node("Add", ["ln", "cst16"], ["Z"], axis=-1),
+                ],
+                "dummy",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT16, ["b", "c", "d"]),
+                    oh.make_tensor_value_info("W", TFLOAT16, ["d"]),
+                    oh.make_tensor_value_info("B", TFLOAT16, ["d"]),
+                ],
+                [oh.make_tensor_value_info("Z", TFLOAT16, ["b", "c", "d"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        dot = to_dot(model)
+        expected = (
+            textwrap.dedent(
+                """
+                digraph {
+                graph [rankdir=TB, splines=true, overlap=false, nodesep=0.2, ranksep=0.2, fontsize=8];
+                node [style="rounded,filled", color="#888888", fontcolor="#222222", shape=box];
+                edge [arrowhead=vee, fontsize=7, labeldistance=-5, labelangle=0];
+                I_0 [label="X\\nFLOAT16(b,c,d)", fillcolor="#aaeeaa"];
+                I_1 [label="W\\nFLOAT16(d)", fillcolor="#aaeeaa"];
+                I_2 [label="B\\nFLOAT16(d)", fillcolor="#aaeeaa"];
+                LayerNormalization_3 [label="LayerNormalization(., ., ., axis=-1)", fillcolor="#cccccc"];
+                Cast_4 [label="Cast([1.0], to=FLOAT16)", fillcolor="#cccccc"];
+                Add_5[label="Add(.,.,axis=-1)",fillcolor="#cccccc"];
+                I_0 -> LayerNormalization_3 [label="FLOAT16(b,c,d)"];
+                I_1 -> LayerNormalization_3 [label="FLOAT16(d)"];
+                I_2 -> LayerNormalization_3 [label="FLOAT16(d)"];
+                LayerNormalization_3 -> Add_5 [label="FLOAT16(b,c,d)"];
+                Cast_4->Add_5[label="FLOAT16()"];
+                O_6 [label="Z\\nFLOAT16(b,c,d)", fillcolor="#aaaaee"];
+                Add_5 -> O_6;
+                }
+                """
+            )
+            .strip("\n")
+            .replace(" ", "")
+        )
+        self.maxDiff = None
+        self.assertEqual(expected, dot.strip("\n").replace(" ", ""))
+
     @requires_transformers("4.57")
     def test_dot_plot_tiny(self):
         data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")

onnx_diagnostic/helpers/dot_helper.py

@@ -27,7 +27,7 @@ def _get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
 
 
 def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
-    els = [f"{node.domain}.\\n{node.op_type}" if node.domain else node.op_type, "("]
+    els = [f"{node.domain}.\\n{node.op_type}" if node.domain else node.op_type, "\\n("]
     ee = [tiny_inits.get(i, ".") if i else "" for i in node.input]
     for att in node.attribute:
         if att.name == "to":
@@ -44,7 +44,10 @@ def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
     els.append(")")
     if node.op_type == "Constant":
         els.extend([" -> ", node.output[0]])
-    return "".join(els)
+    res = "".join(els)
+    if len(res) < 40:
+        return res.replace("\\n(", "(")
+    return res
 
 
 def _make_edge_label(value_info: onnx.ValueInfoProto, multi_line: bool = False) -> str:
@@ -172,7 +175,7 @@ def _mkn(obj: object) -> int:
         inits.append(onh.from_array(value, name=node.output[0]))
 
     for init in inits:
-        if init.name in inputs:
+        if init.name in name_to_ids:
             # hide optional inputs
             continue
         shape = tuple(init.dims)
@@ -188,7 +191,7 @@ def _mkn(obj: object) -> int:
             edge_label[init.name] = ls
 
     for node in nodes:
-        if node.op_type == "Constant" and node.output[0] in name_to_ids:
+        if node.op_type == "Constant" and node.output[0] in tiny_inits:
             continue
         color = op_type_colors.get(node.op_type, "#cccccc")
         label = _make_node_label(node, tiny_inits)