| 275 | } |
| 276 | |
| 277 | std::pair<std::vector<mx::array>, nb::object> tree_flatten_with_structure( |
| 278 | nb::object tree, |
| 279 | bool strict /* = true */) { |
| 280 | auto sentinel = structure_sentinel(); |
| 281 | std::vector<mx::array> flat_tree; |
| 282 | auto structure = tree_map( |
| 283 | tree, |
| 284 | [&flat_tree, sentinel = std::move(sentinel), strict](nb::handle obj) { |
| 285 | if (nb::isinstance<mx::array>(obj)) { |
| 286 | flat_tree.push_back(nb::cast<mx::array>(obj)); |
| 287 | return sentinel; |
| 288 | } else if (!strict) { |
| 289 | return nb::cast<nb::object>(obj); |
| 290 | } else { |
| 291 | throw std::invalid_argument( |
| 292 | "[tree_flatten] The argument should contain only arrays"); |
| 293 | } |
| 294 | }); |
| 295 | |
| 296 | return {flat_tree, structure}; |
| 297 | } |
| 298 | |
| 299 | nb::object tree_unflatten_from_structure( |
| 300 | nb::object structure, |
no test coverage detected