Remove unnecessary argument.

dmlc · Aug 7, 2024 · 93a519e · 93a519e
1 parent 9c09380
commit 93a519e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 10 deletions.
diff --git a/graphbolt/include/graphbolt/cuda_sampling_ops.h b/graphbolt/include/graphbolt/cuda_sampling_ops.h
@@ -56,9 +56,6 @@ namespace ops {
  * sampled edges, see arXiv:2210.13339.
  * @param return_eids Boolean indicating whether edge IDs need to be returned,
  * typically used when edge features are required.
- * @param indices_return_is_optional Boolean indicating whether returning the
- * output indices is optional, typically used if the fetch of indices will
- * happen later.
  * @param type_per_edge A tensor representing the type of each edge, if present.
  * @param probs_or_mask An optional tensor with (unnormalized) probabilities
  * corresponding to each neighboring edge of a node. It must be
@@ -81,7 +78,7 @@ c10::intrusive_ptr<sampling::FusedSampledSubgraph> SampleNeighbors(
     torch::optional<torch::Tensor> seeds,
     torch::optional<std::vector<int64_t>> seed_offsets,
     const std::vector<int64_t>& fanouts, bool replace, bool layer,
-    bool return_eids, bool indices_return_is_optional,
+    bool return_eids,
     torch::optional<torch::Tensor> type_per_edge = torch::nullopt,
     torch::optional<torch::Tensor> probs_or_mask = torch::nullopt,
     torch::optional<torch::Tensor> node_type_offset = torch::nullopt,

diff --git a/graphbolt/src/cuda/neighbor_sampler.cu b/graphbolt/src/cuda/neighbor_sampler.cu
@@ -202,8 +202,7 @@ c10::intrusive_ptr<sampling::FusedSampledSubgraph> SampleNeighbors(
     torch::optional<torch::Tensor> seeds,
     torch::optional<std::vector<int64_t>> seed_offsets,
     const std::vector<int64_t>& fanouts, bool replace, bool layer,
-    bool return_eids, bool indices_return_is_optional,
-    torch::optional<torch::Tensor> type_per_edge,
+    bool return_eids, torch::optional<torch::Tensor> type_per_edge,
     torch::optional<torch::Tensor> probs_or_mask,
     torch::optional<torch::Tensor> node_type_offset,
     torch::optional<torch::Dict<std::string, int64_t>> node_type_to_id,
@@ -520,7 +519,9 @@ c10::intrusive_ptr<sampling::FusedSampledSubgraph> SampleNeighbors(
           }
         }
 
-        if (!indices_return_is_optional || utils::is_on_gpu(indices)) {
+        // TODO @mfbalin: remove true from here once fetching indices later is
+        // setup.
+        if (true || layer || utils::is_on_gpu(indices)) {
           output_indices = Gather(indices, picked_eids);
         }
       }));

diff --git a/graphbolt/src/fused_csc_sampling_graph.cc b/graphbolt/src/fused_csc_sampling_graph.cc
@@ -836,9 +836,9 @@ c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
         c10::DeviceType::CUDA, "SampleNeighbors", {
           return ops::SampleNeighbors(
               indptr_, indices_, seeds, seed_offsets, fanouts, replace, layer,
-              return_eids, false, type_per_edge_, probs_or_mask,
-              node_type_offset_, node_type_to_id_, edge_type_to_id_,
-              random_seed, seed2_contribution);
+              return_eids, type_per_edge_, probs_or_mask, node_type_offset_,
+              node_type_to_id_, edge_type_to_id_, random_seed,
+              seed2_contribution);
         });
   }
   TORCH_CHECK(seeds.has_value(), "Nodes can not be None on the CPU.");