MCPcopy Create free account
hub / github.com/ml-explore/mlx / print_graph

Function print_graph

mlx/graph_utils.cpp:62–103  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

60}
61
62void 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
105void export_to_dot(
106 std::ostream& os,

Callers

nothing calls this directly

Calls 3

depth_first_traversalFunction · 0.85
push_backMethod · 0.80
nameMethod · 0.45

Tested by

no test coverage detected