def construct_subgraph(G, nodes):
subg = GraphView()
nodes = set(nodes)
for node in nodes:
for edge in G.out_edges(node.name):
# only add internal edges
if edge.to_node in nodes:
subg.add_edge(NNEdge(from_node=edge.from_node, to_node=edge.to_node,
from_idx=edge.from_idx, to_idx=edge.to_idx))
def red_fn(state, edge):
state.setdefault((edge.from_node, edge.from_idx), []
).append((edge.to_node, edge.to_idx))
return state
inputs = reduce(red_fn, set([edge for node in subg.inputs()
for edge in G.in_edges(node.name)]), {})
inputs_map = []
for (fnode, fidx), outs in inputs.items():
inp = FusionInputParameters(
f'{fnode.name}_{fidx}_in', dims=fnode.out_dims[fidx])
inputs_map.append((fnode, fidx))
for (tnode, tidx) in outs:
subg.add_edge(NNEdge(from_node=inp, to_node=tnode, to_idx=tidx))
outputs = [(node, set(edge.from_idx for edge in G.out_edges(node.name)))
for node in subg.outputs()]
outputs_map = []
for (node, fidxes) in outputs:
output_map = []
outputs_map.append(output_map)
for fidx in fidxes:
output_map.append((edge.to_node, edge.to_idx)
for edge in G.out_edges(node.name) if edge.from_idx == fidx)
outp = FusionOutputParameters(
f'{node.name}_{fidx}_out', dims=node.out_dims[fidx])
subg.add_edge(NNEdge(from_node=node, to_node=outp, from_idx=fidx))
return (subg, inputs_map, outputs_map)
def new_eval_copy(subg: GraphView) -> CopySet:
node_slices = {}
for inp in subg.inputs():
walk_down(subg, inp, node_slices)
return CopySet([(node, node_slices[node][0]) for node in subg.outputs()])
