Function tree_replace

python/src/trees.cpp:219–234

Replace all the arrays from the src values with the dst values in the tree

// Replace all the arrays from the src values with the dst values in the tree
void tree_replace(
    nb::object& tree,
    const std::vector<mx::array>& src,
    const std::vector<mx::array>& dst) {
  std::unordered_map<uintptr_t, mx::array> src_to_dst;
  for (int i = 0; i < src.size(); ++i) {
    src_to_dst.insert({src[i].id(), dst[i]});
  }
  tree_visit_update(tree, [&](nb::handle node) {
    auto arr = nb::cast<mx::array>(node);
    if (auto it = src_to_dst.find(arr.id()); it != src_to_dst.end()) {
      return nb::cast(it->second);
    }
    return nb::cast(arr);
  });
}

std::vector<mx::array> tree_flatten(nb::handle tree, bool strict /* = true */) {
  std::vector<mx::array> flat_tree;