Minor refactor for temporal sampling (#257)

* disable undirected for temporal sampling

* disjoint sampling for temporal

* fix repeated node index

* compile fix

* Update csrc/cpu/neighbor_sample_cpu.cpp

Co-authored-by: Zecheng Zhang <[email protected]>

* Update csrc/cpu/neighbor_sample_cpu.cpp

Co-authored-by: Zecheng Zhang <[email protected]>

* Update csrc/cpu/neighbor_sample_cpu.cpp

Co-authored-by: Zecheng Zhang <[email protected]>

* comments on directed to be true

* add directed in API

* comments

* minor function signature fix

* Update csrc/cpu/neighbor_sample_cpu.cpp

* Update csrc/neighbor_sample.cpp

* minor refactor

* minor refactor

Co-authored-by: Zecheng Zhang <[email protected]>
Co-authored-by: Matthias Fey <[email protected]>

Author: 3 people

csrc/cpu/neighbor_sample_cpu.cpp

@@ -242,10 +242,11 @@ hetero_sample(const vector<node_t> &node_types,
               // note that the sampling always needs to have directed=True
               // for temporal case
               // to_local_src_node is not used for temporal / directed case
+              const int64_t sample_idx = src_samples.size();
               src_samples.push_back(v);
               src_root_time.push_back(dst_time);
               cols.push_back(i);
-              rows.push_back(src_samples.size() - 1);
+              rows.push_back(sample_idx);
               edges.push_back(offset);
             } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
@@ -271,10 +272,11 @@ hetero_sample(const vector<node_t> &node_types,
               // force disjoint of computation tree
               // note that the sampling always needs to have directed=True
               // for temporal case
+              const int64_t sample_idx = src_samples.size();
               src_samples.push_back(v);
               src_root_time.push_back(dst_time);
               cols.push_back(i);
-              rows.push_back(src_samples.size() - 1);
+              rows.push_back(sample_idx);
               edges.push_back(offset);
             } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
@@ -305,10 +307,11 @@ hetero_sample(const vector<node_t> &node_types,
               // force disjoint of computation tree
               // note that the sampling always needs to have directed=True
               // for temporal case
+              const int64_t sample_idx = src_samples.size();
               src_samples.push_back(v);
               src_root_time.push_back(dst_time);
               cols.push_back(i);
-              rows.push_back(src_samples.size() - 1);
+              rows.push_back(sample_idx);
               edges.push_back(offset);
             } else {
               const auto res = to_local_src_node.insert({v, src_samples.size()});
@@ -431,7 +434,7 @@ hetero_temporal_neighbor_sample_cpu(
     const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
     const c10::Dict<node_t, torch::Tensor> &node_time_dict,
     const int64_t num_hops, const bool replace, const bool directed) {
-  AT_ASSERTM(directed, "Temporal sampling requires 'directed' sampling")
+  AT_ASSERTM(directed, "Temporal sampling requires 'directed' sampling");
   if (replace) {
     // We assume that directed = True for temporal sampling
     // The current implementation uses disjoint computation trees