[GitHub] [incubator-tvm] broune commented on a change in pull request #4370: [WIP] Relay visualization: exporter + visualizer

[GitBox <gi...@apache.org>]

broune commented on a change in pull request #4370: [WIP] Relay visualization: exporter + visualizer
URL: https://github.com/apache/incubator-tvm/pull/4370#discussion_r348139803
 
 

 ##########
 File path: python/tvm/relay/analysis.py
 ##########
 @@ -408,3 +411,116 @@ def structural_hash(value):
         msg = ("found value of type {0} expected" +
                "relay.Expr or relay.Type").format(type(value))
         raise TypeError(msg)
+
+def _export_as_relayviz(expr):
+    """Export a Relay function as a nested dictionary, following the RelayViz spec
+    (https://discuss.tvm.ai/t/rfc-visualizing-relay-program-as-graph/4825/10). The dictionary will
+    contain all information useful for visualizing the Relay program and is meant to be consumed
+    by other visualizers.
+
+    Parameters
+    ----------
+    expr : tvm.relay.Expr
+        The input expression.
+
+    Returns
+    -------
+    viz : dict
+        Nested dictionary
+    """
+
+    # node_dict maps a Relay node to an index (node ID)
+    def _traverse_expr(node, node_dict):
+        if node in node_dict:
+            return
+        node_dict[node] = len(node_dict)
+
+    node_dict = {}
+    post_order_visit(expr, lambda x: _traverse_expr(x, node_dict))
+
+    relayviz_nodes = []
+
+    # Sort by node ID
+    for node, node_idx in sorted(node_dict.items(), key=lambda x: x[1]):
+        if isinstance(node, Function):
+            relayviz_nodes.append({
+                'node_kind': 'Function',
+                'body': node_dict[node.body],
+                'params': [node_dict[x] for x in node.params],
+                'ret_type': {
+                    'dtype': node.ret_type.dtype,
+                    'shape': [int(x) for x in node.ret_type.shape]
+                }
 
 Review comment:
   How about there is a type attribute on every op and it contains a string that describes the type of what the node yields? It can be an empty string if there is no value. Any visualizer that wants to actually understand the structure of e.g. types can just deal with Relay directly.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services