TensorMap CPUHeteroInducer::InitNode()

in graphlearn_torch/csrc/cpu/inducer.cc [93:114]

• 22 lines of code
• 4 McCabe index (conditional complexity)

TensorMap CPUHeteroInducer::InitNode(const TensorMap& seed) {
  Reset();
  TensorMap nodes_dict;
  for (const auto& iter :seed) {
    const auto seed_size = iter.second.size(0);
    const auto seed_ptr = iter.second.data_ptr<int64_t>();
    std::vector<int64_t> out_nodes;
    out_nodes.reserve(seed_size);
    int64_t n_id = nodes_size_[iter.first];
    for (int32_t i = 0; i < seed_size; ++i) {
      if (((glob2local_[iter.first]).insert({seed_ptr[i], n_id})).second) {
        out_nodes.push_back(seed_ptr[i]);
        ++n_id;
      }
    }
    nodes_size_[iter.first] = n_id;
    torch::Tensor nodes = torch::empty(out_nodes.size(), iter.second.options());
    std::copy(out_nodes.begin(), out_nodes.end(), nodes.data_ptr<int64_t>());
    nodes_dict.emplace(iter.first, std::move(nodes));
  }
  return nodes_dict;
}