def __init__(self,
nodes,
*,
duplicate_constants=True,
duplicate_free_variables=True,
function_in_node=True,
tooltip_gen=None,
class_gen=None,
extra_style=None):
super().__init__({'layout': {
'name': 'dagre',
'rankDir': 'TB'
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style)
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.labeler = NodeLabeler(function_in_node=function_in_node,
relation_symbols=short_relation_symbols)
self._class_gen = _make_class_gen(class_gen)
self.todo = set(nodes)
self.graphs = {node.graph for node in nodes if node.graph}
self.focus = set()
# Nodes that are to be colored as return nodes
self.returns = {
node
for node in nodes if node.graph and node is node.graph.return_
}
# IDs for duplicated constants
self.currid = 0
def __init__(self):
"""Initialize."""
self._prim_graph_cache = {}
self.make_const = PythonConstantConverter()
self.graph_to_name = {}
self.fn_name_to_code = {}
self.node_labeler = NodeLabeler(
relation_symbols={"copy": "", "opt": ""}
)
self.name_counter = Counter()
def __init__(
self,
entry_points,
*,
duplicate_constants=False,
duplicate_free_variables=False,
function_in_node=False,
follow_references=False,
tooltip_gen=None,
class_gen=None,
extra_style=None,
beautify=True,
):
"""Initialize a MyiaGraphPrinter."""
super().__init__(
{"layout": {
"name": "dagre",
"rankDir": "TB"
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style,
)
# Graphs left to process
if beautify:
self.graphs = set()
self.focus = set()
for g in entry_points:
self._import_graph(g)
else:
self.graphs = set(entry_points)
self.focus = set(self.graphs)
self.beautify = beautify
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.follow_references = follow_references
self.labeler = NodeLabeler(
function_in_node=function_in_node,
relation_symbols=short_relation_symbols,
)
self._class_gen = _make_class_gen(class_gen)
# Nodes processed
self.processed = set()
# Nodes left to process
self.pool = set()
# Nodes that are to be colored as return nodes
self.returns = set()
# IDs for duplicated constants
self.currid = 0
def __init__(self,
entry_points,
*,
duplicate_constants=False,
duplicate_free_variables=False,
function_in_node=False,
follow_references=False,
tooltip_gen=None,
class_gen=None,
extra_style=None):
"""Initialize a MyiaGraphPrinter."""
super().__init__({'layout': {
'name': 'dagre',
'rankDir': 'TB'
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style)
# Graphs left to process
self.graphs = set(entry_points)
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.follow_references = follow_references
self.labeler = NodeLabeler(function_in_node=function_in_node,
relation_symbols=short_relation_symbols)
self._class_gen = class_gen
# Nodes processed
self.processed = set()
# Nodes left to process
self.pool = set()
# Nodes that are to be colored as return nodes
self.returns = set()
# IDs for duplicated constants
self.currid = 0
# Custom rules for nodes that represent certain calls
self.custom_rules = {
'return': self.process_node_return,
'getitem': self.process_node_getitem,
'make_tuple': self.process_node_make_tuple
}
class PythonCompiler(_Compiler):
"""Compile a myia graph into a Python function."""
REG_PYTHON_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
REG_INVALID_CHARS = re.compile(r"[^A-Za-z0-9_]+")
REG_INVALID_START = re.compile(r"^[^A-Za-z0-9_]+")
REG_INVALID_END = re.compile(r"[^A-Za-z0-9_]+$")
REG_DIGIT_START = re.compile(r"^[0-9]")
def __init__(self):
"""Initialize."""
self._prim_graph_cache = {}
self.make_const = PythonConstantConverter()
self.graph_to_name = {}
self.fn_name_to_code = {}
self.node_labeler = NodeLabeler(relation_symbols={
"copy": "",
"opt": ""
})
self.name_counter = Counter()
def is_valid_python_name(self, text):
"""Return True if text is a valid Python name."""
return self.REG_PYTHON_NAME.match(text)
def convert_to_valid_python_name(self, label):
"""Convert given string to a valid Python variable name."""
formatted_label = label
formatted_label = self.REG_INVALID_START.sub("", formatted_label)
formatted_label = self.REG_INVALID_END.sub("", formatted_label)
formatted_label = self.REG_INVALID_CHARS.sub("_", formatted_label)
if self.REG_DIGIT_START.match(formatted_label):
formatted_label = f"_{formatted_label}"
if not formatted_label:
formatted_label = "var_"
return formatted_label
def get_label(self, node):
"""Generate a label for given node."""
label = self.node_labeler.label(node)
if not self.is_valid_python_name(label):
label = self.convert_to_valid_python_name(label)
assert self.is_valid_python_name(label), label
if label[0] != "_":
label = f"_{label}"
self.name_counter.update([label])
count = self.name_counter[label]
return label if count == 1 else f"{label}_v{count}"
def run(self, graph, backend):
"""Compile given graph.
:type backend: PythonBackend
"""
mng = manage(graph)
mng.keep_roots(graph)
# Graph to name
for g in mng.graphs:
if g is graph:
self.graph_to_name[g] = "main"
else:
self.graph_to_name[g] = self.get_label(g)
# Graph name to function code
for g, g_name in self.graph_to_name.items():
self.fn_name_to_code[g_name] = self.convert_func(g)
# Compilation.
pre_code = [
"import math",
"import operator",
"import numpy as np",
"from myia.utils import RandomStateWrapper",
"from myia.lib import TaggedValue",
"from myia.utils.universe import HandleInstance",
"import myia.compile.backends.python.implementations as IMPL",
]
other_functions = []
main_body = None
main_signature = None
for fn_name, (fn_params, fn_body) in self.fn_name_to_code.items():
if fn_name == "main":
main_body = fn_body
main_signature = f"def main({', '.join(fn_params)}):"
else:
fn_signature = f"def {fn_name}({', '.join(fn_params)}):"
other_functions.append(fn_signature)
other_functions.append(fn_body)
final_structure = (pre_code + other_functions +
[main_signature, main_body])
final_code = nested_list_to_code_string(final_structure)
if backend.debug:
backend.debug.write(f"\n{final_code}")
if backend.pdb:
return PdbRunCall(final_code)
# Compile code string to a Python executable function
# reference: https://stackoverflow.com/a/19850183
compiled = compile(final_code, "", "exec")
module = ModuleType("mod")
exec(compiled, module.__dict__)
return getattr(module, "main")
def has_node(self, node):
"""Return True if given node has already an associated name."""
return isinstance(node, Graph) and node in self.graph_to_name
def has_constant(self, name):
"""Return True if a constant with given name is already registered."""
return name in self.fn_name_to_code
def get_new_name(self, desired_name):
"""Generate a new name usable as a node name."""
if desired_name[0] != "_":
desired_name = f"_{desired_name}"
self.name_counter.update([desired_name])
count = self.name_counter[desired_name]
return desired_name if count == 1 else f"{desired_name}_v{count}"
def ref(self, node):
"""Return a name (string) associated to given node."""
assert node.is_constant_graph() and node.value.parent is None
return self.graph_to_name[node.value]
def get_graph_cache(self):
"""Return graph cache (for graph compilation caching)."""
return self._prim_graph_cache
def convert_func(self, graph):
"""Convert a graph to Python function code."""
return FunctionCompiler(graph, self).compile()
class MyiaNodesPrinter(GraphPrinter):
def __init__(self,
nodes,
*,
duplicate_constants=True,
duplicate_free_variables=True,
function_in_node=True,
tooltip_gen=None,
class_gen=None,
extra_style=None):
super().__init__({'layout': {
'name': 'dagre',
'rankDir': 'TB'
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style)
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.labeler = NodeLabeler(function_in_node=function_in_node,
relation_symbols=short_relation_symbols)
self._class_gen = _make_class_gen(class_gen)
self.todo = set(nodes)
self.graphs = {node.graph for node in nodes if node.graph}
self.focus = set()
# Nodes that are to be colored as return nodes
self.returns = {
node
for node in nodes if node.graph and node is node.graph.return_
}
# IDs for duplicated constants
self.currid = 0
def name(self, x):
"""Return the name of a node."""
return self.labeler.name(x, force=True)
def label(self, node, fn_label=None):
"""Return the label to give to a node."""
return self.labeler.label(node, None, fn_label=fn_label)
def const_fn(self, node):
"""
Return name of function, if constant.
Given an `Apply` node of a constant function, return the
name of that function, otherwise return None.
"""
return self.labeler.const_fn(node)
def add_graph(self, g):
"""Create a node for a graph."""
name = self.name(g)
argnames = [self.name(p) for p in g.parameters]
lbl = f'{name}({", ".join(argnames)})'
classes = ['function', 'focus' if g in self.focus else '']
self.cynode(id=g, label=lbl, classes=' '.join(classes))
# self.processed.add(g)
def process_node_generic(self, node, g, cl):
"""Create node and edges for a node."""
if node.is_constant() and self.duplicate_constants:
return
lbl = self.label(node)
self.cynode(id=node, label=lbl, parent=g, classes=cl)
fn = node.inputs[0] if node.inputs else None
if fn and fn.is_constant_graph():
self.graphs.add(fn.value)
for inp in node.inputs:
if inp.is_constant_graph():
self.cyedge(src_id=g,
dest_id=inp.value,
label=('', 'use-edge'))
edges = []
if fn and not (fn.is_constant() and self.function_in_node):
edges.append((node, 'F', fn))
edges += [(node, i + 1, inp)
for i, inp in enumerate(node.inputs[1:]) or []]
self.process_edges(edges)
def class_gen(self, node, cl=None):
"""Generate the class name for this node."""
g = node.graph
if cl is not None:
pass
elif node in self.returns:
cl = 'output'
elif node.is_parameter():
cl = 'input'
if node not in g.parameters:
cl += ' unlisted'
elif node.is_constant():
cl = 'constant'
elif node.is_special():
cl = f'special-{type(node.special).__name__}'
else:
cl = 'intermediate'
if _has_error(node.debug):
cl += ' error'
if self._class_gen:
return self._class_gen(self._strip_cosmetic(node), cl)
else:
return cl
def process_node(self, node):
"""Create node and edges for a node."""
# if node in self.processed:
# return
g = node.graph
# self.follow(node)
cl = self.class_gen(node)
if node.inputs and node.inputs[0].is_constant():
fn = node.inputs[0].value
if fn in cosmetics:
cosmetics[fn](self, node, g, cl)
elif hasattr(fn, 'graph_display'):
fn.graph_display(self, node, g, cl)
else:
self.process_node_generic(node, g, cl)
else:
self.process_node_generic(node, g, cl)
def process_edges(self, edges):
"""Create edges."""
for edge in edges:
src, lbl, dest = edge
if dest not in self.todo:
continue
if dest.is_constant() and self.duplicate_constants:
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.label(dest),
classes=self.class_gen(dest, 'constant'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
elif self.duplicate_free_variables and \
src.graph and dest.graph and \
src.graph is not dest.graph:
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.labeler.label(dest, force=True),
classes=self.class_gen(dest, 'freevar'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
self.cyedge(src_id=cid, dest_id=dest, label=(lbl, 'link-edge'))
self.cyedge(src_id=src.graph,
dest_id=dest.graph,
label=('', 'nest-edge'))
else:
self.cyedge(src_id=src, dest_id=dest, label=lbl)
def process(self):
"""Process all graphs in entry_points."""
if self.nodes or self.edges:
return
for g in self.graphs:
self.add_graph(g)
for node in self.todo:
self.process_node(node)
return self.nodes, self.edges
class MyiaGraphPrinter(GraphPrinter):
"""
Utility to generate a graphical representation for a graph.
Attributes:
duplicate_constants: Whether to create a separate node for
every instance of the same constant.
duplicate_free_variables: Whether to create a separate node
to represent the use of a free variable, or point directly
to that node in a different graph.
function_in_node: Whether to print, when possible, the name
of a node's operation directly in the node's label instead
of creating a node for the operation and drawing an edge
to it.
follow_references: Whether to also print graphs that are
called by this graph.
"""
def __init__(self,
entry_points,
*,
duplicate_constants=False,
duplicate_free_variables=False,
function_in_node=False,
follow_references=False,
tooltip_gen=None,
class_gen=None,
extra_style=None,
beautify=True):
"""Initialize a MyiaGraphPrinter."""
super().__init__({'layout': {
'name': 'dagre',
'rankDir': 'TB'
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style)
# Graphs left to process
if beautify:
self.graphs = set()
self.focus = set()
for g in entry_points:
self._import_graph(g)
else:
self.graphs = set(entry_points)
self.focus = set(self.graphs)
self.beautify = beautify
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.follow_references = follow_references
self.labeler = NodeLabeler(function_in_node=function_in_node,
relation_symbols=short_relation_symbols)
self._class_gen = _make_class_gen(class_gen)
# Nodes processed
self.processed = set()
# Nodes left to process
self.pool = set()
# Nodes that are to be colored as return nodes
self.returns = set()
# IDs for duplicated constants
self.currid = 0
def _import_graph(self, graph):
mng = manage(graph, weak=True)
graphs = set()
parents = mng.parents
g = graph
while g:
graphs.add(g)
g = parents[g]
clone = GraphCloner(*graphs, total=True, relation='cosmetic')
self.graphs |= {clone[g] for g in graphs}
self.focus.add(clone[graph])
def name(self, x):
"""Return the name of a node."""
return self.labeler.name(x, force=True)
def label(self, node, fn_label=None):
"""Return the label to give to a node."""
return self.labeler.label(node, None, fn_label=fn_label)
def const_fn(self, node):
"""
Return name of function, if constant.
Given an `Apply` node of a constant function, return the
name of that function, otherwise return None.
"""
return self.labeler.const_fn(node)
def add_graph(self, g):
"""Create a node for a graph."""
if g in self.processed:
return
if self.beautify:
g = cosmetic_transformer(g)
name = self.name(g)
argnames = [self.name(p) for p in g.parameters]
lbl = f'{name}({", ".join(argnames)})'
classes = ['function', 'focus' if g in self.focus else '']
self.cynode(id=g, label=lbl, classes=' '.join(classes))
self.processed.add(g)
def process_node_generic(self, node, g, cl):
"""Create node and edges for a node."""
lbl = self.label(node)
self.cynode(id=node, label=lbl, parent=g, classes=cl)
fn = node.inputs[0] if node.inputs else None
if fn and fn.is_constant_graph():
self.graphs.add(fn.value)
for inp in node.inputs:
if inp.is_constant_graph():
self.cyedge(src_id=g,
dest_id=inp.value,
label=('', 'use-edge'))
edges = []
if fn and not (fn.is_constant() and self.function_in_node):
edges.append((node, 'F', fn))
edges += [(node, i + 1, inp)
for i, inp in enumerate(node.inputs[1:]) or []]
self.process_edges(edges)
def class_gen(self, node, cl=None):
"""Generate the class name for this node."""
g = node.graph
if cl is not None:
pass
elif node in self.returns:
cl = 'output'
elif node.is_parameter():
cl = 'input'
if node not in g.parameters:
cl += ' unlisted'
elif node.is_constant():
cl = 'constant'
elif node.is_special():
cl = f'special-{type(node.special).__name__}'
else:
cl = 'intermediate'
if _has_error(node.debug):
cl += ' error'
if self._class_gen:
return self._class_gen(self._strip_cosmetic(node), cl)
else:
return cl
def process_node(self, node):
"""Create node and edges for a node."""
if node in self.processed:
return
g = node.graph
self.follow(node)
cl = self.class_gen(node)
if g and g not in self.processed:
self.add_graph(g)
if node.inputs and node.inputs[0].is_constant():
fn = node.inputs[0].value
if fn in cosmetics:
cosmetics[fn](self, node, g, cl)
elif hasattr(fn, 'graph_display'):
fn.graph_display(self, node, g, cl)
else:
self.process_node_generic(node, g, cl)
else:
self.process_node_generic(node, g, cl)
self.processed.add(node)
def process_edges(self, edges):
"""Create edges."""
for edge in edges:
src, lbl, dest = edge
if dest.is_constant() and self.duplicate_constants:
self.follow(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.label(dest),
classes=self.class_gen(dest, 'constant'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
elif self.duplicate_free_variables and \
src.graph and dest.graph and \
src.graph is not dest.graph:
self.pool.add(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.labeler.label(dest, force=True),
classes=self.class_gen(dest, 'freevar'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
self.cyedge(src_id=cid, dest_id=dest, label=(lbl, 'link-edge'))
self.cyedge(src_id=src.graph,
dest_id=dest.graph,
label=('', 'nest-edge'))
else:
self.pool.add(dest)
self.cyedge(src_id=src, dest_id=dest, label=lbl)
def process_graph(self, g):
"""Process a graph."""
self.add_graph(g)
for inp in g.parameters:
self.process_node(inp)
if not g.return_:
return
ret = g.return_.inputs[1]
if not ret.is_apply() or ret.graph is not g:
ret = g.return_
self.returns.add(ret)
self.pool.add(ret)
while self.pool:
node = self.pool.pop()
self.process_node(node)
def process(self):
"""Process all graphs in entry_points."""
if self.nodes or self.edges:
return
while self.graphs:
g = self.graphs.pop()
self.process_graph(g)
return self.nodes, self.edges
def follow(self, node):
"""Add this node's graph if follow_references is True."""
if node.is_constant_graph() and self.follow_references:
self.graphs.add(node.value)
class MyiaGraphPrinter(GraphPrinter):
"""
Utility to generate a graphical representation for a graph.
Attributes:
duplicate_constants: Whether to create a separate node for
every instance of the same constant.
duplicate_free_variables: Whether to create a separate node
to represent the use of a free variable, or point directly
to that node in a different graph.
function_in_node: Whether to print, when possible, the name
of a node's operation directly in the node's label instead
of creating a node for the operation and drawing an edge
to it.
follow_references: Whether to also print graphs that are
called by this graph.
"""
def __init__(self,
entry_points,
*,
duplicate_constants=False,
duplicate_free_variables=False,
function_in_node=False,
follow_references=False,
tooltip_gen=None,
class_gen=None,
extra_style=None):
"""Initialize a MyiaGraphPrinter."""
super().__init__({'layout': {
'name': 'dagre',
'rankDir': 'TB'
}},
tooltip_gen=tooltip_gen,
extra_style=extra_style)
# Graphs left to process
self.graphs = set(entry_points)
self.duplicate_constants = duplicate_constants
self.duplicate_free_variables = duplicate_free_variables
self.function_in_node = function_in_node
self.follow_references = follow_references
self.labeler = NodeLabeler(function_in_node=function_in_node,
relation_symbols=short_relation_symbols)
self._class_gen = class_gen
# Nodes processed
self.processed = set()
# Nodes left to process
self.pool = set()
# Nodes that are to be colored as return nodes
self.returns = set()
# IDs for duplicated constants
self.currid = 0
# Custom rules for nodes that represent certain calls
self.custom_rules = {
'return': self.process_node_return,
'getitem': self.process_node_getitem,
'make_tuple': self.process_node_make_tuple
}
def name(self, x):
"""Return the name of a node."""
return self.labeler.name(x, force=True)
def label(self, node, fn_label=None):
"""Return the label to give to a node."""
return self.labeler.label(node, None, fn_label=fn_label)
def const_fn(self, node):
"""
Return name of function, if constant.
Given an `Apply` node of a constant function, return the
name of that function, otherwise return None.
"""
return self.labeler.const_fn(node)
def add_graph(self, g):
"""Create a node for a graph."""
if g in self.processed:
return
name = self.name(g)
argnames = [self.name(p) for p in g.parameters]
lbl = f'{name}({", ".join(argnames)})'
self.cynode(id=g, label=lbl, classes='function')
self.processed.add(g)
def process_node_return(self, node, g, cl):
"""Create node and edges for `return ...`."""
self.cynode(id=node, label='', parent=g, classes='const_output')
ret = node.inputs[1]
self.process_edges([(node, '', ret)])
def process_node_getitem(self, node, g, cl):
"""Create node and edges for `x[ct]`."""
idx = node.inputs[2]
if self.function_in_node and is_constant(idx):
lbl = self.label(node, '')
self.cynode(id=node, label=lbl, parent=g, classes=cl)
self.process_edges([(node, (f'[{idx.value}]', 'fn-edge'),
node.inputs[1])])
else:
self.process_node_generic(node, g, cl)
def process_node_make_tuple(self, node, g, cl):
"""Create node and edges for `(a, b, c, ...)`."""
if self.function_in_node:
lbl = self.label(node, f'(...)')
self.cynode(id=node, label=lbl, parent=g, classes=cl)
edges = [(node, i + 1, inp)
for i, inp in enumerate(node.inputs[1:]) or []]
self.process_edges(edges)
else:
return self.process_node_generic(node, g, cl)
def process_node_generic(self, node, g, cl):
"""Create node and edges for a node."""
lbl = self.label(node)
self.cynode(id=node, label=lbl, parent=g, classes=cl)
fn = node.inputs[0] if node.inputs else None
if fn and is_constant_graph(fn):
self.graphs.add(fn.value)
edges = []
if fn and not (is_constant(fn) and self.function_in_node):
edges.append((node, 'F', fn))
edges += [(node, i + 1, inp)
for i, inp in enumerate(node.inputs[1:]) or []]
self.process_edges(edges)
def class_gen(self, node, cl=None):
"""Generate the class name for this node."""
g = node.graph
if cl is not None:
pass
elif node in self.returns:
cl = 'output'
elif g and node in g.parameters:
cl = 'input'
elif is_constant(node):
cl = 'constant'
else:
cl = 'intermediate'
if self._class_gen:
return self._class_gen(node, cl)
else:
return cl
def process_node(self, node):
"""Create node and edges for a node."""
if node in self.processed:
return
g = node.graph
self.follow(node)
cl = self.class_gen(node)
ctfn = self.const_fn(node)
if ctfn:
if ctfn in self.custom_rules:
self.custom_rules[ctfn](node, g, cl)
else:
self.process_node_generic(node, g, cl)
else:
self.process_node_generic(node, g, cl)
self.processed.add(node)
def process_edges(self, edges):
"""Create edges."""
for edge in edges:
src, lbl, dest = edge
if is_constant(dest) and self.duplicate_constants:
self.follow(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.label(dest),
classes=self.class_gen(dest, 'constant'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
elif self.duplicate_free_variables and \
src.graph and dest.graph and \
src.graph is not dest.graph:
self.pool.add(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.name(dest),
classes=self.class_gen(dest, 'freevar'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
self.cyedge(src_id=cid, dest_id=dest, label=(lbl, 'link-edge'))
else:
self.pool.add(dest)
self.cyedge(src_id=src, dest_id=dest, label=lbl)
def process_graph(self, g):
"""Process a graph."""
self.add_graph(g)
for inp in g.parameters:
self.process_node(inp)
if not g.return_:
return
ret = g.return_.inputs[1]
if not is_apply(ret) or ret.graph is not g:
ret = g.return_
self.returns.add(ret)
self.pool.add(ret)
while self.pool:
node = self.pool.pop()
self.process_node(node)
def process(self):
"""Process all graphs in entry_points."""
if self.nodes or self.edges:
return
while self.graphs:
g = self.graphs.pop()
self.process_graph(g)
return self.nodes, self.edges
def follow(self, node):
"""Add this node's graph if follow_references is True."""
if is_constant_graph(node) and self.follow_references:
self.graphs.add(node.value)
