add edge labels (#678)

src/ngraph/pass/visualize_tree.cpp

@@ -19,6 +19,7 @@
 #include "ngraph/function.hpp"
 #include "ngraph/graph_util.hpp"
 #include "ngraph/node.hpp"
+#include "ngraph/ops/get_output_element.hpp"
 #include "ngraph/pass/pass.hpp"
 #include "ngraph/pass/visualize_tree.hpp"
 #include "ngraph/util.hpp"
@@ -32,12 +33,27 @@ bool pass::VisualizeTree::run_on_module(vector<shared_ptr<ngraph::Function>>& fu
     {
         // map<size_t, list<node_ptr>> dependent_nodes;
         traverse_nodes(f, [&](shared_ptr<Node> node) {
+            size_t i = 0;
             for (auto arg : node->get_input_ops())
             {
                 m_ss << add_attributes(arg);
                 m_ss << add_attributes(node);
                 m_ss << "    " << arg->get_name() << " -> " << node->get_name();
+
+                if (std::getenv("NGRAPH_VISUALIZE_EDGE_LABELS") != nullptr)
+                {
+                    size_t output = 0;
+                    if (auto goe = std::dynamic_pointer_cast<op::GetOutputElement>(node))
+                    {
+                        output = goe->get_n();
+                    }
+                    stringstream label_edge;
+                    label_edge << "[label=\" " << output << " -> " << i << " \"]";
+                    m_ss << label_edge.str();
+                }
+
                 m_ss << ";\n";
+                i++;
             }
         });
     }