| 234 | } |
| 235 | |
| 236 | std::vector<mx::array> tree_flatten(nb::handle tree, bool strict /* = true */) { |
| 237 | std::vector<mx::array> flat_tree; |
| 238 | |
| 239 | tree_visit(tree, [&](nb::handle obj) { |
| 240 | if (nb::isinstance<mx::array>(obj)) { |
| 241 | flat_tree.push_back(nb::cast<mx::array>(obj)); |
| 242 | } else if (strict) { |
| 243 | throw std::invalid_argument( |
| 244 | "[tree_flatten] The argument should contain only arrays"); |
| 245 | } |
| 246 | }); |
| 247 | |
| 248 | return flat_tree; |
| 249 | } |
| 250 | |
| 251 | nb::object tree_unflatten( |
| 252 | nb::object tree, |
no test coverage detected