From d6b62341b10e41a2ab24c519eb528cac9212fa05 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 12:30:02 -0700
Subject: [PATCH] undo changes to uniform neighbor sample

---
 cpp/src/c_api/neighbor_sampling.hpp               |  1 -
 cpp/src/sampling/neighbor_sampling_impl.hpp       |  9 +++------
 cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp |  8 ++------
 cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp |  9 ++-------
 cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp | 10 ++--------
 cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp | 10 ++--------
 cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp | 10 ++--------
 cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp | 10 ++--------
 8 files changed, 15 insertions(+), 52 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
index 9c82f191f3a..9defb8a39dd 100644
--- a/cpp/src/c_api/neighbor_sampling.hpp
+++ b/cpp/src/c_api/neighbor_sampling.hpp
@@ -17,7 +17,6 @@
 #pragma once
 
 #include "c_api/array.hpp"
-// FIXME: Remove this file
 namespace cugraph {
 namespace c_api {
 
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 9adab8ad7ed..fe9cceb1fd1 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -388,10 +388,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<label_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
-  // raft::host_span<int32_t const> fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -410,8 +407,8 @@ uniform_neighbor_sample(
     starting_vertices,
     starting_vertex_labels,
     label_to_output_comm_rank,
-    fan_out,
-    heterogeneous_fan_out,
+    std::make_optional(fan_out),
+    std::nullopt,
     return_hops,
     with_replacement,
     prior_sources_behavior,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
index 1a3e56df26a..5aa325c8578 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
index 0d3526e202f..913e7d8a62e 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -96,7 +92,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
index a5844fecf80..77ab2e476ac 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -96,7 +92,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
@@ -124,7 +119,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
index c82ff6c2c03..05ccdfca1f7 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -96,7 +92,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
@@ -124,7 +119,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
index ae51c392f5c..178b1322d89 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -96,7 +92,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
@@ -124,7 +119,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
index 704313c67fe..ba1edd3f630 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
@@ -39,9 +39,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,9 +65,7 @@ uniform_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  std::optional<raft::host_span<int32_t const>> fan_out,
-  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
-    heterogeneous_fan_out,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -96,7 +92,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
@@ -124,7 +119,6 @@ biased_neighbor_sample(
   std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  // std::optional<raft::host_span<int32_t const>> fan_out,
   raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,