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

Function tree_visit

python/src/trees.cpp:94–151  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

92}
93
94void tree_visit(
95 const std::vector<nb::object>& trees,
96 std::function<void(const std::vector<nb::object>&)> visitor) {
97 std::function<void(const std::vector<nb::object>&)> recurse;
98
99 recurse = [&](const std::vector<nb::object>& subtrees) {
100 if (nb::isinstance<nb::list>(subtrees[0])) {
101 std::vector<nb::object> items(subtrees.size());
102 validate_subtrees<nb::list, nb::tuple, nb::dict>(subtrees);
103 for (int i = 0; i < nb::cast<nb::list>(subtrees[0]).size(); ++i) {
104 for (int j = 0; j < subtrees.size(); ++j) {
105 if (nb::isinstance<nb::list>(subtrees[j])) {
106 items[j] = nb::cast<nb::list>(subtrees[j])[i];
107 } else {
108 items[j] = subtrees[j];
109 }
110 }
111 recurse(items);
112 }
113 } else if (nb::isinstance<nb::tuple>(subtrees[0])) {
114 // Check the rest of the subtrees
115 std::vector<nb::object> items(subtrees.size());
116 int len = nb::cast<nb::tuple>(subtrees[0]).size();
117 validate_subtrees<nb::tuple, nb::list, nb::dict>(subtrees);
118 for (int i = 0; i < len; ++i) {
119 for (int j = 0; j < subtrees.size(); ++j) {
120 if (nb::isinstance<nb::tuple>(subtrees[j])) {
121 items[j] = nb::cast<nb::tuple>(subtrees[j])[i];
122 } else {
123 items[j] = subtrees[j];
124 }
125 }
126 recurse(items);
127 }
128 } else if (nb::isinstance<nb::dict>(subtrees[0])) {
129 std::vector<nb::object> items(subtrees.size());
130 validate_subtrees<nb::dict, nb::list, nb::tuple>(subtrees);
131 for (auto item : nb::cast<nb::dict>(subtrees[0])) {
132 for (int j = 0; j < subtrees.size(); ++j) {
133 if (nb::isinstance<nb::dict>(subtrees[j])) {
134 auto subdict = nb::cast<nb::dict>(subtrees[j]);
135 if (!subdict.contains(item.first)) {
136 throw std::invalid_argument(
137 "[tree_visit] Tree is not a valid prefix tree of the first tree.");
138 }
139 items[j] = subdict[item.first];
140 } else {
141 items[j] = subtrees[j];
142 }
143 }
144 recurse(items);
145 }
146 } else {
147 visitor(subtrees);
148 }
149 };
150 return recurse(trees);
151}

Callers 4

py_value_and_gradFunction · 0.85
py_vmapFunction · 0.85
operator()Method · 0.85
tree_flattenFunction · 0.85

Calls 1

sizeMethod · 0.45

Tested by

no test coverage detected