| 1366 | } |
| 1367 | |
| 1368 | VarNodeArray dump( |
| 1369 | std::shared_ptr<ComputingGraph> graph, |
| 1370 | std::vector<std::tuple<std::string, std::string, TensorShape>> inputs, |
| 1371 | std::vector<std::pair<std::string, std::string>> outputs, |
| 1372 | bool prefer_input_names) { |
| 1373 | auto& self = *this; |
| 1374 | mgb_assert(self.trace_result); |
| 1375 | // mark is like "arg_0", "kwarg_xxx", "output_0" ... |
| 1376 | std::unordered_map<std::string, size_t> mark2var; |
| 1377 | for (size_t i = 0; i < self.trace_result->vars.size(); ++i) { |
| 1378 | auto& name = self.trace_result->vars[i].mark; |
| 1379 | if (!name.empty()) { |
| 1380 | mark2var[name] = i; |
| 1381 | } |
| 1382 | } |
| 1383 | std::vector<std::tuple<size_t, std::string, TensorShape>> input_vars; |
| 1384 | std::vector<std::pair<size_t, std::string>> output_vars; |
| 1385 | for (auto&& [input_mark, input_name, input_shape] : inputs) { |
| 1386 | mgb_assert(input_shape.ndim, "input shape invalid"); |
| 1387 | input_vars.push_back( |
| 1388 | {mark2var.at(input_mark), input_name, input_shape}); |
| 1389 | } |
| 1390 | for (auto&& [output_name, repr] : outputs) { |
| 1391 | output_vars.push_back({mark2var.at(output_name), repr}); |
| 1392 | } |
| 1393 | self.options_visitor(py::cast(&graph->options())); |
| 1394 | auto vars = self.trace_result->dump( |
| 1395 | *graph, input_vars, output_vars, prefer_input_names); |
| 1396 | return vars; |
| 1397 | } |
| 1398 | }; |
| 1399 | |
| 1400 | py::class_<Trace>(m, "Trace") |