def _build_chunk_dag(node_str, edge_str): from mars.tensor.random import TensorRandint from mars.tensor.arithmetic import TensorTreeAdd char_dag = DAG() for s in node_str.split(','): char_dag.add_node(s.strip()) for s in edge_str.split(','): l, r = s.split('->') char_dag.add_edge(l.strip(), r.strip()) chunk_dag = DAG() str_to_chunk = dict() for s in char_dag.topological_iter(): if char_dag.count_predecessors(s): c = TensorTreeAdd(_key=s, dtype=np.float32()).new_chunk(None, shape=(10, 10)) inputs = c.op._inputs = [ str_to_chunk[ps] for ps in char_dag.predecessors(s) ] else: c = TensorRandint(_key=s, dtype=np.float32()).new_chunk(None, shape=(10, 10)) inputs = [] str_to_chunk[s] = c chunk_dag.add_node(c) for inp in inputs: chunk_dag.add_edge(inp, c) return chunk_dag, str_to_chunk
def testCompose(self): """ test compose in build graph and optimize """ r""" graph(@: node, #: composed_node): @ --> @ --> @ ========> # """ chunks = [ TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None) for n in range(3) ] graph = DAG() list(map(graph.add_node, chunks[:3])) graph.add_edge(chunks[0], chunks[1]) graph.add_edge(chunks[1], chunks[2]) composed_nodes = graph.compose() self.assertTrue(composed_nodes[0].composed == chunks[:3]) # make the middle one as result chunk, thus the graph cannot be composed composed_nodes = graph.compose(keys=[chunks[1].key]) self.assertEqual(len(composed_nodes), 0) r""" graph(@: node, #: composed_node): @ @ @ @ \ / \ / @ --> @ ========> # / \ / \ @ @ @ @ """ chunks = [ TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None) for n in range(6) ] graph = DAG() list(map(graph.add_node, chunks[:6])) chunks[2].op._inputs = [chunks[0], chunks[1]] chunks[3].op._inputs = [chunks[2]] chunks[4].op._inputs = [chunks[3]] chunks[5].op._inputs = [chunks[3]] graph.add_edge(chunks[0], chunks[2]) graph.add_edge(chunks[1], chunks[2]) graph.add_edge(chunks[2], chunks[3]) graph.add_edge(chunks[3], chunks[4]) graph.add_edge(chunks[3], chunks[5]) composed_nodes = graph.compose() self.assertTrue(composed_nodes[0].composed == chunks[2:4]) # to make sure the predecessors and successors of compose are right # 0 and 1's successors must be composed self.assertIn(composed_nodes[0], graph.successors(chunks[0])) self.assertIn(composed_nodes[0], graph.successors(chunks[1])) # check composed's inputs self.assertIn(chunks[0].data, composed_nodes[0].inputs) self.assertIn(chunks[1].data, composed_nodes[0].inputs) # check composed's predecessors self.assertIn(chunks[0], graph.predecessors(composed_nodes[0])) self.assertIn(chunks[1], graph.predecessors(composed_nodes[0])) # check 4 and 5's inputs self.assertIn(composed_nodes[0].data, graph.successors(composed_nodes[0])[0].inputs) self.assertIn(composed_nodes[0].data, graph.successors(composed_nodes[0])[0].inputs) # check 4 and 5's predecessors self.assertIn(composed_nodes[0], graph.predecessors(chunks[4])) self.assertIn(composed_nodes[0], graph.predecessors(chunks[5])) # test optimizer compose r""" graph(@: node, S: Slice Chunk, #: composed_node): @ @ @ @ \ / \ / @ --> @ --> S ========> # --> S / \ / \ @ @ @ @ compose stopped at S, because numexpr don't support Slice op """ chunks = [ TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None) for n in range(6) ] chunk_slice = TensorSlice().new_chunk([None], None) graph = DAG() list(map(graph.add_node, chunks[:6])) graph.add_node(chunk_slice) graph.add_edge(chunks[0], chunks[2]) graph.add_edge(chunks[1], chunks[2]) graph.add_edge(chunks[2], chunks[3]) graph.add_edge(chunks[3], chunk_slice) graph.add_edge(chunk_slice, chunks[4]) graph.add_edge(chunk_slice, chunks[5]) optimizer = NeRuntimeOptimizer(graph) composed_nodes = optimizer.compose() self.assertTrue(composed_nodes[0].composed == chunks[2:4]) r""" graph(@: node, S: Slice Chunk, #: composed_node): @ --> @ --> S --> @ ========> # --> S --> @ compose stopped at S, because numexpr don't support Slice op """ chunks = [ TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None) for n in range(4) ] graph = DAG() list(map(graph.add_node, chunks[:3])) graph.add_node(chunk_slice) graph.add_edge(chunks[0], chunks[1]) graph.add_edge(chunks[1], chunk_slice) graph.add_edge(chunk_slice, chunks[2]) optimizer = NeRuntimeOptimizer(graph) composed_nodes = optimizer.compose() self.assertTrue(composed_nodes[0].composed == chunks[:2]) self.assertTrue(len(composed_nodes) == 1) r""" graph(@: node, S: Slice Chunk, #: composed_node): @ --> @ --> S --> @ --> @ ========> # --> S --> # compose stopped at S, because numexpr don't support Slice op """ chunks = [ TensorTreeAdd(args=[], _key=str(n)).new_chunk(None, None) for n in range(4) ] graph = DAG() list(map(graph.add_node, chunks[:4])) graph.add_node(chunk_slice) graph.add_edge(chunks[0], chunks[1]) graph.add_edge(chunks[1], chunk_slice) graph.add_edge(chunk_slice, chunks[2]) graph.add_edge(chunks[2], chunks[3]) optimizer = NeRuntimeOptimizer(graph) composed_nodes = optimizer.compose() self.assertTrue(composed_nodes[0].composed == chunks[:2]) self.assertTrue(composed_nodes[1].composed == chunks[2:4])
def testDAG(self): r""" 1 --- 4 2 --- 6 \ / 5 / 3 """ dag = DAG() [dag.add_node(i) for i in range(1, 7)] dag.add_edge(1, 4) dag.add_edge(2, 6) dag.add_edge(2, 5) dag.add_edge(5, 6) dag.add_edge(3, 5) with self.assertRaises(KeyError): dag.add_edge(1, 10) with self.assertRaises(KeyError): dag.add_edge(10, 1) self.assertEqual(set(dag[2]), {5, 6}) self.assertEqual(list(dag.topological_iter()), [3, 2, 5, 6, 1, 4]) self.assertEqual(list(dag.dfs()), [3, 2, 5, 6, 1, 4]) self.assertEqual(list(dag.bfs()), [1, 2, 3, 4, 5, 6]) dag.add_edge(6, 1) dag.add_edge(1, 2) with self.assertRaises(KeyError): for _ in dag.iter_predecessors(-1): pass with self.assertRaises(KeyError): for _ in dag.iter_successors(-1): pass self.assertRaises(GraphContainsCycleError, lambda: list(dag.topological_iter())) dag.remove_edge(2, 5) self.assertFalse(dag.has_successor(2, 5)) with self.assertRaises(KeyError): dag.remove_edge(2, 5) rev_dag = dag.build_reversed() for n in dag: self.assertIn(n, rev_dag) self.assertTrue( all( rev_dag.has_successor(n, pred) for pred in dag.predecessors(n))) undigraph = dag.build_undirected() for n in dag: self.assertIn(n, undigraph) self.assertTrue( all( undigraph.has_predecessor(pred, n) for pred in dag.predecessors(n))) self.assertTrue( all( undigraph.has_successor(n, pred) for pred in dag.predecessors(n))) dag_copy = dag.copy() for n in dag: self.assertIn(n, dag_copy) self.assertTrue( all( dag_copy.has_successor(pred, n) for pred in dag_copy.predecessors(n)))