def node(self, viz_node):
# add the node first
super().node(viz_node)
# if it's AwesomeVar, duplicate it.
if viz_node.type_name == "AwesomeVar":
duplicated_id = f"duplciated_{viz_node.identity}"
duplicated_type = "double AwesomeVar"
super().node(VizNode(duplicated_id, duplicated_type, ""))
# connect the duplicated var to the original one
super().edge(VizEdge(duplicated_id, viz_node.identity))
def get_node_edges(
self,
node: relay.Expr,
relay_param: Dict[str, tvm.runtime.NDArray],
node_to_id: Dict[relay.Expr, str],
) -> Tuple[Union[VizNode, None], List[VizEdge]]:
if isinstance(node, relay.Var):
node = VizNode(node_to_id[node], "AwesomeVar", f"name_hint {node.name_hint}")
# no edge is introduced. So return an empty list.
return node, []
# delegate other types to the other parser.
return self._delegate.get_node_edges(node, relay_param, node_to_id)