yobx.helpers.mermaid_helper#

yobx.helpers.mermaid_helper.to_mermaid(model: ModelProto) → str[source]#

Converts a model into a Mermaid flowchart TD string.

The function:

  • uses BasicShapeBuilder to annotate every edge with its inferred dtype and shape (when available),

  • inlines small scalar constants and 1-D initializers directly onto the node label so the graph stays compact,

  • handles Scan / Loop / If sub-graphs by drawing dotted edges for outer-scope values consumed by the sub-graph.

Parameters:

model – ONNX model to convert

Returns:

Mermaid flowchart string

<<<

import numpy as np
import onnx
import onnx.helper as oh
import onnx.numpy_helper as onh
from yobx.helpers.mermaid_helper import to_mermaid

TFLOAT = onnx.TensorProto.FLOAT
model = oh.make_model(
    oh.make_graph(
        [
            oh.make_node("Add", ["X", "Y"], ["added"]),
            oh.make_node("MatMul", ["added", "W"], ["mm"]),
            oh.make_node("Relu", ["mm"], ["Z"]),
        ],
        "add_matmul_relu",
        [
            oh.make_tensor_value_info("X", TFLOAT, ["batch", "seq", 4]),
            oh.make_tensor_value_info("Y", TFLOAT, ["batch", "seq", 4]),
        ],
        [oh.make_tensor_value_info("Z", TFLOAT, ["batch", "seq", 2])],
        [onh.from_array(np.zeros((4, 2), dtype=np.float32), name="W")],
    ),
    opset_imports=[oh.make_opsetid("", 18)],
    ir_version=10,
)
mermaid_src = to_mermaid(model)
print(mermaid_src)

>>>

    flowchart TD
        I_0["X\nFLOAT(batch,seq,4)"]:::input
        I_1["Y\nFLOAT(batch,seq,4)"]:::input
        i_2["W\nFLOAT(4, 2)"]:::init
        Add_3["Add(., .)"]:::op
        MatMul_4["MatMul(., .)"]:::op
        Relu_5["Relu(.)"]:::op
        I_0 -->|"FLOAT(batch,seq,4)"| Add_3
        I_1 -->|"FLOAT(batch,seq,4)"| Add_3
        Add_3 -->|"FLOAT(batch,seq,4)"| MatMul_4
        i_2 -->|"FLOAT(4, 2)"| MatMul_4
        MatMul_4 -->|"FLOAT(batch,seq,2)"| Relu_5
        O_6["Z\nFLOAT(batch,seq,2)"]:::output
        Relu_5 --> O_6
        classDef input fill:#aaeeaa,stroke:#00aa00,color:#000
        classDef init fill:#cccc00,stroke:#888800,color:#000
        classDef op fill:#cccccc,stroke:#666666,color:#000
        classDef output fill:#aaaaee,stroke:#0000aa,color:#000

flowchart TD
    I_0["X\nFLOAT(batch,seq,4)"]:::input
    I_1["Y\nFLOAT(batch,seq,4)"]:::input
    i_2["W\nFLOAT(4, 2)"]:::init
    Add_3["Add(., .)"]:::op
    MatMul_4["MatMul(., .)"]:::op
    Relu_5["Relu(.)"]:::op
    I_0 -->|"FLOAT(batch,seq,4)"| Add_3
    I_1 -->|"FLOAT(batch,seq,4)"| Add_3
    Add_3 -->|"FLOAT(batch,seq,4)"| MatMul_4
    i_2 -->|"FLOAT(4, 2)"| MatMul_4
    MatMul_4 -->|"FLOAT(batch,seq,2)"| Relu_5
    O_6["Z\nFLOAT(batch,seq,2)"]:::output
    Relu_5 --> O_6
    classDef input fill:#aaeeaa,stroke:#00aa00,color:#000
    classDef init fill:#cccc00,stroke:#888800,color:#000
    classDef op fill:#cccccc,stroke:#666666,color:#000
    classDef output fill:#aaaaee,stroke:#0000aa,color:#000