exnviz.py 文件源码

def begin_pass(self, filename=None, **kwargs):
        super(ExVizPass, self).begin_pass(**kwargs)
        try:
            import graphviz
        except ImportError:
            raise ImportError("You tried to use the ShowGraph transformer pass but did "
                              "not have the python graphviz library installed")
        if filename is None:
            filename = self.filename
        self.exops_with_nodes = set()
        self.exops_without_nodes = set()
        self.tensors_with_nodes = set()
        self.tensors_without_nodes = set()

        # Get all ops from this set
        self.graph = graphviz.Digraph(name=filename,
                                      # node_attr={'shape': 'box', 'style': 'rounded'},
                                      graph_attr={'nodesep': '.5',
                                                  'ranksep': '.5'})