SubGraph CPUSubGraphOp::NodeSubGraph()

in graphlearn_torch/csrc/cpu/subgraph_op.cc [20:44]

• 23 lines of code
• 2 McCabe index (conditional complexity)

SubGraph CPUSubGraphOp::NodeSubGraph(const torch::Tensor& srcs,
                                     bool with_edge) {
  std::vector<int64_t> out_nodes;
  InitNode(srcs, out_nodes);
  torch::Tensor nodes = torch::empty(out_nodes.size(), srcs.options());
  std::copy(out_nodes.begin(), out_nodes.end(), nodes.data_ptr<int64_t>());
  std::vector<int64_t> out_rows;
  std::vector<int64_t> out_cols;
  std::vector<int64_t> out_eids;
  Induce(out_nodes, with_edge, out_rows, out_cols, out_eids);

  auto edge_size = out_rows.size();
  torch::Tensor rows = torch::empty(edge_size, srcs.options());
  torch::Tensor cols = torch::empty(edge_size, srcs.options());
  std::copy(out_rows.begin(), out_rows.end(), rows.data_ptr<int64_t>());
  std::copy(out_cols.begin(), out_cols.end(), cols.data_ptr<int64_t>());
  auto subgraph = SubGraph(nodes, rows, cols);
  if (with_edge) {
    torch::Tensor eids = torch::empty(edge_size, srcs.options());
    std::copy(out_eids.begin(), out_eids.end(), eids.data_ptr<int64_t>());
    subgraph.eids = eids;
  }

  return subgraph;
}