| 60 | } |
| 61 | |
| 62 | void print_graph( |
| 63 | std::ostream& os, |
| 64 | NodeNamer namer, |
| 65 | const std::vector<array>& outputs) { |
| 66 | std::vector<array> tape; |
| 67 | std::vector<array> inputs; |
| 68 | |
| 69 | depth_first_traversal( |
| 70 | [&](const array& x) { |
| 71 | if (x.has_primitive()) { |
| 72 | tape.push_back(x); |
| 73 | } else { |
| 74 | inputs.push_back(x); |
| 75 | } |
| 76 | }, |
| 77 | outputs); |
| 78 | |
| 79 | auto print_arrs = [&namer, &os](std::vector<array> arrs) { |
| 80 | for (auto& arr : arrs) { |
| 81 | os << namer.get_name(arr); |
| 82 | os << " [" << arr.shape() << ", " << arr.dtype() << "]"; |
| 83 | if (&arr != &arrs.back()) { |
| 84 | os << ", "; |
| 85 | } |
| 86 | } |
| 87 | }; |
| 88 | |
| 89 | os << "Inputs: "; |
| 90 | print_arrs(inputs); |
| 91 | os << "\nOutputs: "; |
| 92 | print_arrs(outputs); |
| 93 | os << "\n"; |
| 94 | |
| 95 | for (auto& arr : tape) { |
| 96 | os << arr.primitive().name(); |
| 97 | os << " "; |
| 98 | print_arrs(arr.inputs()); |
| 99 | os << " -> "; |
| 100 | print_arrs(arr.outputs()); |
| 101 | os << "\n"; |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | void export_to_dot( |
| 106 | std::ostream& os, |
nothing calls this directly
no test coverage detected