From 0adb2fdb31e3dc24c00df069348bb8017020aef8 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 13 Aug 2024 08:49:42 -0700
Subject: [PATCH 001/117] support heterogenous fanout type

---
 cpp/include/cugraph/sampling_functions.hpp    | 223 +++++++++++++++++-
 cpp/include/cugraph_c/sampling_algorithms.h   |  47 ++++
 cpp/src/c_api/array.hpp                       |  12 +
 cpp/src/c_api/graph_functions.cpp             |   2 +-
 cpp/src/c_api/neighbor_sampling.cpp           | 150 +++++++++++-
 cpp/src/c_api/neighbor_sampling.hpp           |  30 +++
 cpp/src/sampling/neighbor_sampling_impl.hpp   |  47 ++--
 .../sampling/neighbor_sampling_mg_v32_e32.cpp |   8 +-
 .../sampling/neighbor_sampling_mg_v32_e64.cpp |   8 +-
 .../sampling/neighbor_sampling_mg_v64_e64.cpp |   8 +-
 .../sampling/neighbor_sampling_sg_v32_e32.cpp |   8 +-
 .../sampling/neighbor_sampling_sg_v32_e64.cpp |   8 +-
 .../sampling/neighbor_sampling_sg_v64_e64.cpp |   8 +-
 .../dask/sampling/uniform_neighbor_sample.py  |  15 +-
 .../sampling/uniform_neighbor_sample.py       |  17 +-
 .../_cugraph_c/sampling_algorithms.pxd        |  18 ++
 .../pylibcugraph/uniform_neighbor_sample.pyx  |  75 ++++--
 17 files changed, 634 insertions(+), 50 deletions(-)
 create mode 100644 cpp/src/c_api/neighbor_sampling.hpp

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index fec1a07604e..5b884fa611e 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -129,7 +129,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement                           = true,
@@ -137,6 +138,114 @@ uniform_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
+#if 0
+/* FIXME:
+ There are two options to support heterogeneous fanout
+ 1) Create a new function 'heterogeneous_uniform_neighbor_sample' which will take as input
+      only heterogeneous fanout type. Drawback: code redundancy
+ 2) Update 'uniform_neighbor_sample' to support both fanout types
+*/
+/**
+ * @brief Heterogeneous uniform Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
+ * (optional) will be described below and is dependent upon the input parameters.
+ *
+ * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
+ * label and offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
+ * the label output has values.  This will also result in the output being sorted by vertex label.
+ * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
+ * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
+ * label output has values.  This will also result in the output being sorted by vertex label.  The
+ * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
+ * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
+ * be shuffled so that all data with a particular label will be on the specified rank.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam label_t Type of label. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * for the sampling.
+ * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
+ * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
+ * output rank.  The label span must be sorted in ascending order.
+ * @param fan_out Tuple of host spans mapping each edge type to fanout values. Element 0
+ * of the tuple defines the size of the fanout per edge type while element 1 defines the branching out
+ * (fan-out) degree per edge type for each level.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param return_hops boolean flag specifying if the hop information should be returned
+ * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
+ * DEFAULT)
+ * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
+ * multiple times with the same label, it will only be passed once (for each label) as a source
+ * for the next hop.  Default is false.
+ * @param with_replacement boolean flag specifying if random sampling is done with replacement
+ * (true); or, without replacement (false); default = true;
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+
+// tuple with 3 elements. 1 - edge_type, 2- host span of size_t, 3 - fanout vector 
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>> fan_out,
+  raft::random::RngState& rng_state,
+  bool return_hops,
+  bool with_replacement                           = true,
+  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
+  bool dedupe_sources                             = false,
+  bool do_expensive_check                         = false);
+#endif
+
 /**
  * @brief Biased Neighborhood Sampling.
  *
@@ -240,6 +349,118 @@ biased_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
+#if 0
+/* FIXME:
+ There are two options to support heterogeneous fanout
+ 1) Create a new function 'heterogeneous_biased_neighbor_sample' which will take as input
+      only heterogeneous fanout type. Drawback: code redundancy
+ 2) Update 'biased_neighbor_sample' to support both fanout types
+*/
+/**
+ * @brief Biased Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
+ * (optional) will be described below and is dependent upon the input parameters.
+ *
+ * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
+ * label and offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
+ * the label output has values.  This will also result in the output being sorted by vertex label.
+ * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
+ * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
+ * label output has values.  This will also result in the output being sorted by vertex label.  The
+ * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
+ * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
+ * be shuffled so that all data with a particular label will be on the specified rank.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam label_t Type of label. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
+ * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
+ * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
+ * corresponding edge can never be selected.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * for the sampling.
+ * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
+ * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
+ * output rank.  The label span must be sorted in ascending order.
+ * @param fan_out Tuple of host spans mapping each edge type to fanout values. Element 0
+ * of the tuple defines the size of the fanout per edge type while element 1 defines the branching out
+ * (fan-out) degree per edge type for each level.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param return_hops boolean flag specifying if the hop information should be returned
+ * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
+ * DEFAULT)
+ * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
+ * multiple times with the same label, it will only be passed once (for each label) as a source
+ * for the next hop.  Default is false.
+ * @param with_replacement boolean flag specifying if random sampling is done with replacement
+ * (true); or, without replacement (false); default = true;
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>> fan_out,
+  raft::random::RngState& rng_state,
+  bool return_hops,
+  bool with_replacement                           = true,
+  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
+  bool dedupe_sources                             = false,
+  bool do_expensive_check                         = false);
+#endif
+
 /*
  * @brief renumber sampled edge list and compress to the (D)CSR|(D)CSC format.
  *
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 1a3d20b9339..e440512c482 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -319,6 +319,49 @@ void cugraph_sampling_set_dedupe_sources(cugraph_sampling_options_t* options, bo
  */
 void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
 
+/**
+ * @brief     Opaque neighborhood sampling heterogeneous fanout type
+ */
+// FIXME: internal representation should be tuple instead of pairs - Make it more generic (tuple)
+// cugraph_device_tuple_t, host_device_tuple_t,
+//dictionary, key and array
+// translate dictionary to a tuple. Add to the draft PR the PLC layer.
+// Concatenate to build the 3 arrays from the PLC layer
+/// mimic 
+typedef struct {
+  int32_t align_;
+} cugraph_sample_heterogeneous_fanout_t;
+
+/**
+ * @brief       Create heterogeneous fanout
+ *
+ * Input data will be stored in the heterogenous_fanout.
+ *
+ * @param [in]  handle         Handle for accessing resources
+ * @param [in]  graph          Pointer to graph
+ * @param [in]  edge_type_size Type erased array of edge type size
+ * @param [in]  fanout         Type erased array of fanout values
+ * @param [out] heterogeneous_fanout Opaque pointer to fanout_t
+ * @param [out] error          Pointer to an error object storing details of any error.  Will
+ *                             be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_create_heterogeneous_fanout(
+  const cugraph_resource_handle_t* handle,
+  cugraph_graph_t* graph,
+  const cugraph_type_erased_host_array_view_t* edge_type_size,
+  const cugraph_type_erased_host_array_view_t* fanout,
+  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+  cugraph_error_t** error);
+
+/**
+ * @brief     Free edge type and fanout pairs
+ *
+ * @param [in]    heterogeneous_fanout The edge type size and fanout values
+ */
+void cugraph_heterogeneous_fanout_free(
+  cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
+
 /**
  * @brief     Uniform Neighborhood Sampling
  *
@@ -368,6 +411,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -667,6 +711,8 @@ cugraph_error_code_t cugraph_test_uniform_neighborhood_sample_result_create(
  *                                not CUGRAPH_SUCCESS
  * @return error code
  */
+
+
 cugraph_error_code_t cugraph_select_random_vertices(const cugraph_resource_handle_t* handle,
                                                     const cugraph_graph_t* graph,
                                                     cugraph_rng_state_t* rng_state,
@@ -674,6 +720,7 @@ cugraph_error_code_t cugraph_select_random_vertices(const cugraph_resource_handl
                                                     cugraph_type_erased_device_array_t** vertices,
                                                     cugraph_error_t** error);
 
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/cpp/src/c_api/array.hpp b/cpp/src/c_api/array.hpp
index 048d2ee1cea..2dce70fa924 100644
--- a/cpp/src/c_api/array.hpp
+++ b/cpp/src/c_api/array.hpp
@@ -125,6 +125,18 @@ struct cugraph_type_erased_host_array_t {
     std::copy(vec.begin(), vec.end(), reinterpret_cast<T*>(data_.get()));
   }
 
+  template <typename T>
+  T* as_type()
+  {
+    return reinterpret_cast<T*>(data_.get());
+  }
+
+  template <typename T>
+  T const* as_type() const
+  {
+    return reinterpret_cast<T const*>(data_.get());
+  }
+
   auto view()
   {
     return new cugraph_type_erased_host_array_view_t{data_.get(), size_, num_bytes_, type_};
diff --git a/cpp/src/c_api/graph_functions.cpp b/cpp/src/c_api/graph_functions.cpp
index df741a349d2..349598b7ef6 100644
--- a/cpp/src/c_api/graph_functions.cpp
+++ b/cpp/src/c_api/graph_functions.cpp
@@ -84,7 +84,7 @@ struct create_vertex_pairs_functor : public cugraph::c_api::abstract_functor {
                               std::nullopt,
                               std::nullopt);
       }
-
+      // std::tuple (template)
       result_ = new cugraph::c_api::cugraph_vertex_pairs_t{
         new cugraph::c_api::cugraph_type_erased_device_array_t(first_copy, graph_->vertex_type_),
         new cugraph::c_api::cugraph_type_erased_device_array_t(second_copy, graph_->vertex_type_)};
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 69306806030..18e158cbd6d 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -20,8 +20,10 @@
 #include "c_api/random.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
+#include "c_api/neighbor_sampling.hpp" // FIXME: Remove this and instead use std::tuple
 
 #include <cugraph_c/algorithms.h>
+#include <cugraph_c/sampling_algorithms.h>
 
 #include <cugraph/algorithms.hpp>
 #include <cugraph/detail/shuffle_wrappers.hpp>
@@ -72,6 +74,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
+  cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out_{nullptr};
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool do_expensive_check_{false};
@@ -86,6 +89,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
     cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
     cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
+    cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out,
     cugraph_rng_state_t* rng_state,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool do_expensive_check)
@@ -108,6 +112,8 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
           label_offsets)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
+      heterogeneous_fan_out_(
+        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       options_(options),
       do_expensive_check_(do_expensive_check)
@@ -215,7 +221,17 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
                 raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
                                                  label_to_comm_rank_->size_}))
             : std::nullopt,
-          raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
+          (fan_out_ != nullptr)
+            ? std::make_optional<raft::host_span<const int>>(fan_out_->as_type<const int>(), fan_out_->size_)
+            : std::nullopt,
+    
+          (heterogeneous_fan_out_ != nullptr)
+            ? std::make_optional(std::make_tuple(
+              raft::host_span<const int>{heterogeneous_fan_out_->edge_type_id->as_type<int>(),
+                                         heterogeneous_fan_out_->edge_type_id->size_},
+              raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
+                                         heterogeneous_fan_out_->fanout->size_}))
+            : std::nullopt,
           rng_state_->rng_state_,
           options_.return_hops_,
           options_.with_replacement_,
@@ -753,6 +769,121 @@ struct biased_neighbor_sampling_functor : public cugraph::c_api::abstract_functo
   }
 };
 
+  struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
+    raft::handle_t const& handle_;
+    cugraph::c_api::cugraph_graph_t* graph_;
+    cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
+    cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
+    // FIXME: This type doesn't exist: instead create an 'std::tuple<cugraph_type_erased_host_array_t*>'
+    cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
+
+    create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
+                                ::cugraph_graph_t* graph,
+                                ::cugraph_type_erased_host_array_view_t const* edge_type_size,
+                                ::cugraph_type_erased_host_array_view_t const* fanout)
+      : abstract_functor(),
+        handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+        graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+        edge_type_size_(
+          reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_size)),
+        fanout_(
+          reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
+    {
+    }
+
+    template <typename vertex_t,
+              typename edge_t,
+              typename weight_t,
+              typename edge_type_type_t,
+              bool store_transposed,
+              bool multi_gpu>
+    void operator()()
+    {
+      // FIXME: Remove this check as it is not necessary
+      if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
+        unsupported();
+      } else {
+        std::vector<int32_t> edge_type_size_copy{(int32_t) edge_type_size_->size_};
+        std::vector<int32_t> fanout_copy{(int32_t) fanout_->size_};
+      
+        raft::copy(
+          edge_type_size_copy.data(), edge_type_size_->as_type<int32_t>(), edge_type_size_->size_, handle_.get_stream());
+      
+        raft::copy(
+          fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
+        
+        // std::tuple (template)
+        // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
+        result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
+          new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, INT32),
+          new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)};
+        
+      }
+    }
+  };
+
+  #if 0
+  struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
+    raft::handle_t const& handle_;
+    //cugraph::c_api::cugraph_graph_t* graph_;
+    //cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
+    //cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
+    // FIXME: This type doesn't exist: instead create an 'std::tuple<cugraph_type_erased_host_array_t*>'
+    //cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
+
+    create_heterogeneous_fanout_functor(cugraph_resource_handle_t const* handle//,
+                                //::cugraph_graph_t* graph,
+                                //::cugraph_type_erased_host_array_view_t const* edge_type_size,
+                                //::cugraph_type_erased_host_array_view_t const* fanout
+                                )
+      : abstract_functor(),
+        handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_)//,
+        //graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+        //edge_type_size_(
+        //  reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_size)),
+        //fanout_(
+        //  reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
+    {
+    }
+
+    template <typename vertex_t,
+              typename edge_t,
+              typename weight_t,
+              typename edge_type_t,
+              bool store_transposed,
+              bool multi_gpu>
+    void operator()()
+    {
+      printf("\nhello world\n");
+      /*
+      // FIXME: Remove this check as it is not necessary
+      if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
+        unsupported();
+      } else {
+        std::vector<int32_t> edge_type_size_copy{edge_type_size_->size_};
+        std::vector<int32_t> fanout_copy{fanout_->size_};
+        //rmm::device_uvector<int32_t> edge_type_size_copy(edge_type_size_->size_, handle_.get_stream());
+        //rmm::device_uvector<int32_t> fanout_copy(fanout_->size_, handle_.get_stream());
+
+        raft::copy(
+          edge_type_size_copy.data(), edge_type_size_->as_type<vertex_t>(), edge_type_size_->size_, handle_.get_stream());
+        raft::copy(
+          fanout_copy.data(), fanout_->as_type<vertex_t>(), fanout_->size_, handle_.get_stream());
+
+        // std::tuple (template)
+        // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
+        result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
+          new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, graph_->vertex_type_),
+          new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, graph_->vertex_type_)};
+      }
+      */
+    }
+  };
+  #endif
+
+
+
+
 }  // namespace
 
 extern "C" cugraph_error_code_t cugraph_sampling_options_create(
@@ -1246,6 +1377,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -1300,6 +1432,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                             label_to_comm_rank,
                                             label_offsets,
                                             fan_out,
+                                            heterogeneous_fanout,
                                             rng_state,
                                             std::move(options_cpp),
                                             do_expensive_check};
@@ -1383,3 +1516,18 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                            do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
+
+
+
+extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
+  const cugraph_resource_handle_t* handle,
+  cugraph_graph_t* graph,
+  const cugraph_type_erased_host_array_view_t* edge_type_size,
+  const cugraph_type_erased_host_array_view_t* fanout,
+  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+  cugraph_error_t** error)
+{
+  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_size, fanout);  
+
+  return cugraph::c_api::run_algorithm(graph, functor, heterogeneous_fanout, error);
+}
\ No newline at end of file
diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
new file mode 100644
index 00000000000..346ad150fab
--- /dev/null
+++ b/cpp/src/c_api/neighbor_sampling.hpp
@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "c_api/array.hpp"
+// FIXME: Remove this file
+namespace cugraph {
+namespace c_api {
+
+struct cugraph_sample_heterogeneous_fanout_t {
+  cugraph_type_erased_host_array_t* edge_type_id;
+  cugraph_type_erased_host_array_t* fanout;
+};
+
+}  // namespace c_api
+}  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 1785f934426..1989b733eaf 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -59,7 +59,9 @@ neighbor_sample_impl(
   std::optional<raft::device_span<label_t const>> this_frontier_vertex_labels,
   std::optional<std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>>
     label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
+  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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -75,11 +77,24 @@ neighbor_sample_impl(
 #else
   static_assert(std::is_floating_point_v<bias_t>);
 
-  CUGRAPH_EXPECTS(fan_out.size() > 0, "Invalid input argument: number of levels must be non-zero.");
-  CUGRAPH_EXPECTS(
-    fan_out.size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
+  if (fan_out) {
+    CUGRAPH_EXPECTS((*fan_out).size() > 0, "Invalid input argument: number of levels must be non-zero.");
+    CUGRAPH_EXPECTS(
+    (*fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
     "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                               // to store hops
+  } else {    
+    CUGRAPH_EXPECTS(std::accumulate(
+      std::get<0>(*heterogeneous_fan_out).begin(),
+      std::get<0>(*heterogeneous_fan_out).end(), 0) == std::get<1>(*heterogeneous_fan_out).size() && std::get<1>(*heterogeneous_fan_out).size() != 0,
+        "Invalid input argument: number of levels and size must match and should be non zero.");
+    
+    CUGRAPH_EXPECTS(
+      std::get<0>(*heterogeneous_fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max())
+        && std::get<1>(*heterogeneous_fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
+      "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
+                                                                                // to store hops
+  }
 
   if constexpr (!multi_gpu) {
     CUGRAPH_EXPECTS(!label_to_output_comm_rank,
@@ -123,12 +138,12 @@ neighbor_sample_impl(
     this_frontier_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
                                 : std::nullopt;
 
-  level_result_src_vectors.reserve(fan_out.size());
-  level_result_dst_vectors.reserve(fan_out.size());
-  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve(fan_out.size()); }
-  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve(fan_out.size()); }
-  if (level_result_edge_type_vectors) { (*level_result_edge_type_vectors).reserve(fan_out.size()); }
-  if (level_result_label_vectors) { (*level_result_label_vectors).reserve(fan_out.size()); }
+  level_result_src_vectors.reserve((*fan_out).size());
+  level_result_dst_vectors.reserve((*fan_out).size());
+  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve((*fan_out).size()); }
+  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve((*fan_out).size()); }
+  if (level_result_edge_type_vectors) { (*level_result_edge_type_vectors).reserve((*fan_out).size()); }
+  if (level_result_label_vectors) { (*level_result_label_vectors).reserve((*fan_out).size()); }
 
   rmm::device_uvector<vertex_t> frontier_vertices(0, handle.get_stream());
   auto frontier_vertex_labels =
@@ -150,7 +165,7 @@ neighbor_sample_impl(
 
   std::vector<size_t> level_sizes{};
   int32_t hop{0};
-  for (auto&& k_level : fan_out) {
+  for (auto&& k_level : (*fan_out)) {
     rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
     rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());
     std::optional<rmm::device_uvector<weight_t>> weights{std::nullopt};
@@ -192,7 +207,7 @@ neighbor_sample_impl(
     if (labels) { (*level_result_label_vectors).push_back(std::move(*labels)); }
 
     ++hop;
-    if (hop < fan_out.size()) {
+    if (hop < (*fan_out).size()) {
       // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span
       //  rather than making a copy.
       auto vertex_partition_range_lasts = graph_view.vertex_partition_range_lasts();
@@ -304,7 +319,7 @@ neighbor_sample_impl(
   if (return_hops) {
     result_hops   = rmm::device_uvector<int32_t>(result_size, handle.get_stream());
     output_offset = 0;
-    for (size_t i = 0; i < fan_out.size(); ++i) {
+    for (size_t i = 0; i < (*fan_out).size(); ++i) {
       scalar_fill(
         handle, result_hops->data() + output_offset, level_sizes[i], static_cast<int32_t>(i));
       output_offset += level_sizes[i];
@@ -366,7 +381,9 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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,
@@ -386,6 +403,7 @@ uniform_neighbor_sample(
     starting_vertex_labels,
     label_to_output_comm_rank,
     fan_out,
+    heterogeneous_fan_out,
     return_hops,
     with_replacement,
     prior_sources_behavior,
@@ -440,6 +458,7 @@ biased_neighbor_sample(
     starting_vertex_labels,
     label_to_output_comm_rank,
     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 f61c1c10c53..04cba83ac53 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
index c37b353ae1c..6ba353ad58c 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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 ea3f6b466da..a33bff7d3b1 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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 0f0affbb323..44f82c0b0ea 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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 7ab0a8782ec..69128c02fd2 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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 70dd9a59842..056730f8100 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
@@ -39,7 +39,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -65,7 +66,8 @@ 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,
-  raft::host_span<int32_t const> fan_out,
+  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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -92,6 +94,7 @@ 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,
@@ -119,6 +122,7 @@ 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/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index 03746561817..688a9d973ef 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from __future__ import annotations
+from itertools import chain
 
 import warnings
 
@@ -377,9 +378,10 @@ def uniform_neighbor_sample(
     start_list : int, list, cudf.Series, or dask_cudf.Series (int32 or int64)
         a list of starting vertices for sampling
 
-    fanout_vals : list
+    fanout_vals : list (int32) or dict
         List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
+        hop level or dictionary of edge type and fanout values for
+        heterogeneous fanout type.
 
     with_replacement: bool, optional (default=True)
         Flag to specify if the random sampling is done with replacement
@@ -612,6 +614,15 @@ def uniform_neighbor_sample(
         fanout_vals = fanout_vals.get().astype("int32")
     elif isinstance(fanout_vals, cudf.Series):
         fanout_vals = fanout_vals.values_host.astype("int32")
+    elif isinstance(fanout_vals, dict):
+        # FIXME: Add expensive check to ensure all dict values are lists
+        # Convert to a tuple of sequence (edge type size and fanout values)
+        edge_type_size = []
+        [edge_type_size.append(len(s)) for s in list(fanout_vals.values())]
+        edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
+        fanout_vals = (
+            numpy.asarray(edge_type_size, dtype="int32"),
+            numpy.asarray(edge_type_fanout_vals, dtype="int32"))
     else:
         raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
 
diff --git a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
index 2e3d032a20f..7ba482f36bd 100644
--- a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
@@ -12,6 +12,7 @@
 # limitations under the License.
 
 from __future__ import annotations
+from itertools import chain
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import uniform_neighbor_sample as pylibcugraph_uniform_neighbor_sample
@@ -90,9 +91,10 @@ def uniform_neighbor_sample(
     start_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
 
-    fanout_vals : list (int32)
+    fanout_vals : list (int32) or dict
         List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
+        hop level or dictionary of edge type and fanout values for
+        heterogeneous fanout type.
 
     with_replacement: bool, optional (default=True)
         Flag to specify if the random sampling is done with replacement
@@ -314,8 +316,17 @@ def uniform_neighbor_sample(
         fanout_vals = fanout_vals.get().astype("int32")
     elif isinstance(fanout_vals, cudf.Series):
         fanout_vals = fanout_vals.values_host.astype("int32")
+    elif isinstance(fanout_vals, dict):
+        # FIXME: Add expensive check to ensure all dict values are lists
+        # Convert to a tuple of sequence (edge type size and fanout values)
+        edge_type_size = []
+        [edge_type_size.append(len(s)) for s in list(fanout_vals.values())]
+        edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
+        fanout_vals = (
+            numpy.asarray(edge_type_size, dtype="int32"),
+            numpy.asarray(edge_type_fanout_vals, dtype="int32"))
     else:
-        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
+        raise TypeError("fanout_vals must be a sequence or a dictionary, " f"got: {type(fanout_vals)}")
 
     if "weights" in G.edgelist.edgelist_df:
         weight_t = G.edgelist.edgelist_df["weights"].dtype
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 0f852d9cecd..a9cea82c3d7 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -56,6 +56,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
         const cugraph_type_erased_host_array_view_t* fan_out,
+        const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
         cugraph_rng_state_t* rng_state,
         const cugraph_sampling_options_t* options,
         bool_t do_expensive_check,
@@ -103,3 +104,20 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
             cugraph_type_erased_device_array_t** vertices,
             cugraph_error_t** error
         )
+    
+    ctypedef struct cugraph_sample_heterogeneous_fanout_t:
+        pass
+    
+    cdef cugraph_error_code_t \
+        cugraph_create_heterogeneous_fanout(
+            const cugraph_resource_handle_t* handle,
+            cugraph_graph_t* graph,
+            const cugraph_type_erased_host_array_view_t* edge_type_size,
+            const cugraph_type_erased_host_array_view_t* fanout,
+            cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+            cugraph_error_t** error
+        )
+
+    cdef void \
+        cugraph_heterogeneous_fanout_free(
+            cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
diff --git a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
index f002622f497..fee88b4d16f 100644
--- a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
@@ -53,7 +53,9 @@ from pylibcugraph._cugraph_c.algorithms cimport (
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
     cugraph_uniform_neighbor_sample,
-
+    cugraph_sample_heterogeneous_fanout_t,
+    cugraph_create_heterogeneous_fanout,
+    cugraph_heterogeneous_fanout_free,
 )
 from pylibcugraph.resource_handle cimport (
     ResourceHandle,
@@ -83,7 +85,6 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
                             _GPUGraph input_graph,
                             start_list,
                             h_fan_out,
-                            *,
                             bool_t with_replacement,
                             bool_t do_expensive_check,
                             with_edge_properties=False,
@@ -116,7 +117,7 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     start_list: device array type
         Device array containing the list of starting vertices for sampling.
 
-    h_fan_out: numpy array type
+    h_fan_out: numpy array type or tuple of numpy array type
         Device array containing the brancing out (fan-out) degrees per
         starting vertex for each hop level.
 
@@ -211,22 +212,70 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     )
 
     cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fanout_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
 
     cdef bool_t c_deduplicate_sources = deduplicate_sources
     cdef bool_t c_return_hops = return_hops
     cdef bool_t c_renumber = renumber
     cdef bool_t c_compress_per_hop = compress_per_hop
 
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+    cdef uintptr_t ai_fan_out_size_ptr
+    cdef uintptr_t ai_fan_out_values_ptr
+
+
     assert_CAI_type(start_list, "start_list")
     assert_CAI_type(batch_id_list, "batch_id_list", True)
     assert_CAI_type(label_list, "label_list", True)
     assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
     assert_CAI_type(label_offsets, "label_offsets", True)
-    assert_AI_type(h_fan_out, "h_fan_out")
+    if not isinstance(h_fan_out, tuple):
+        assert_AI_type(h_fan_out, "h_fan_out")
+        ai_fan_out_ptr = \
+            h_fan_out.__array_interface__["data"][0]
+        
+        fan_out_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_ptr,
+                len(h_fan_out),
+                get_c_type_from_numpy_type(h_fan_out.dtype))
+    else:
+        assert_AI_type(h_fan_out[0], "h_fan_out_size")
+        assert_AI_type(h_fan_out[1], "h_fan_out_values")
+        ai_fan_out_size_ptr = \
+            h_fan_out[0].__array_interface__["data"][0]
+        ai_fan_out_values_ptr = \
+            h_fan_out[1].__array_interface__["data"][0]
+
+        fan_out_size_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_size_ptr,
+                len(h_fan_out[0]),
+                get_c_type_from_numpy_type(h_fan_out[0].dtype))
+    
+        fan_out_values_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_values_ptr,
+                len(h_fan_out[1]),
+                get_c_type_from_numpy_type(h_fan_out[1].dtype))
+        
+        error_code = cugraph_create_heterogeneous_fanout(
+            c_resource_handle_ptr,
+            c_graph_ptr,
+            fan_out_size_ptr,
+            fan_out_values_ptr,
+            &heterogeneous_fanout_ptr,
+            &error_ptr
+        )
+
+    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fanout")
 
     cdef cugraph_sample_result_t* result_ptr
-    cdef cugraph_error_code_t error_code
-    cdef cugraph_error_t* error_ptr
 
     cdef uintptr_t cai_start_ptr = \
         start_list.__cuda_array_interface__["data"][0]
@@ -251,9 +300,6 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
         cai_label_offsets_ptr = \
             label_offsets.__cuda_array_interface__['data'][0]
 
-    cdef uintptr_t ai_fan_out_ptr = \
-        h_fan_out.__array_interface__["data"][0]
-
     cdef cugraph_type_erased_device_array_view_t* start_ptr = \
         cugraph_type_erased_device_array_view_create(
             <void*>cai_start_ptr,
@@ -298,13 +344,6 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
                 get_c_type_from_numpy_type(label_offsets.dtype)
             )
 
-    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = \
-        cugraph_type_erased_host_array_view_create(
-            <void*>ai_fan_out_ptr,
-            len(h_fan_out),
-            get_c_type_from_numpy_type(h_fan_out.dtype))
-
-
     cg_rng_state = CuGraphRandomState(resource_handle, random_state)
 
     cdef cugraph_rng_state_t* rng_state_ptr = \
@@ -362,6 +401,7 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
         label_to_output_comm_rank_ptr,
         label_offsets_ptr,
         fan_out_ptr,
+        heterogeneous_fanout_ptr,
         rng_state_ptr,
         sampling_options,
         do_expensive_check,
@@ -372,6 +412,9 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     # Free the sampling options
     cugraph_sampling_options_free(sampling_options)
 
+    if isinstance(h_fan_out, tuple):
+        cugraph_heterogeneous_fanout_free(heterogeneous_fanout_ptr)
+
     # Free the two input arrays that are no longer needed.
     cugraph_type_erased_device_array_view_free(start_ptr)
     cugraph_type_erased_host_array_view_free(fan_out_ptr)

From bb5a3e2f5ffcde22932f405804a8e538827a832e Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 13 Aug 2024 11:01:13 -0700
Subject: [PATCH 002/117] remove unusued code

---
 cpp/src/c_api/neighbor_sampling.cpp | 63 -----------------------------
 1 file changed, 63 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 18e158cbd6d..9e6e68871c6 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -821,69 +821,6 @@ struct biased_neighbor_sampling_functor : public cugraph::c_api::abstract_functo
       }
     }
   };
-
-  #if 0
-  struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
-    raft::handle_t const& handle_;
-    //cugraph::c_api::cugraph_graph_t* graph_;
-    //cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
-    //cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
-    // FIXME: This type doesn't exist: instead create an 'std::tuple<cugraph_type_erased_host_array_t*>'
-    //cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
-
-    create_heterogeneous_fanout_functor(cugraph_resource_handle_t const* handle//,
-                                //::cugraph_graph_t* graph,
-                                //::cugraph_type_erased_host_array_view_t const* edge_type_size,
-                                //::cugraph_type_erased_host_array_view_t const* fanout
-                                )
-      : abstract_functor(),
-        handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_)//,
-        //graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-        //edge_type_size_(
-        //  reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_size)),
-        //fanout_(
-        //  reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
-    {
-    }
-
-    template <typename vertex_t,
-              typename edge_t,
-              typename weight_t,
-              typename edge_type_t,
-              bool store_transposed,
-              bool multi_gpu>
-    void operator()()
-    {
-      printf("\nhello world\n");
-      /*
-      // FIXME: Remove this check as it is not necessary
-      if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
-        unsupported();
-      } else {
-        std::vector<int32_t> edge_type_size_copy{edge_type_size_->size_};
-        std::vector<int32_t> fanout_copy{fanout_->size_};
-        //rmm::device_uvector<int32_t> edge_type_size_copy(edge_type_size_->size_, handle_.get_stream());
-        //rmm::device_uvector<int32_t> fanout_copy(fanout_->size_, handle_.get_stream());
-
-        raft::copy(
-          edge_type_size_copy.data(), edge_type_size_->as_type<vertex_t>(), edge_type_size_->size_, handle_.get_stream());
-        raft::copy(
-          fanout_copy.data(), fanout_->as_type<vertex_t>(), fanout_->size_, handle_.get_stream());
-
-        // std::tuple (template)
-        // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
-        result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
-          new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, graph_->vertex_type_),
-          new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, graph_->vertex_type_)};
-      }
-      */
-    }
-  };
-  #endif
-
-
-
-
 }  // namespace
 
 extern "C" cugraph_error_code_t cugraph_sampling_options_create(

From 10fa86d17e5520c046ee0ba1144d28c70f19e39f Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 13 Aug 2024 12:22:19 -0700
Subject: [PATCH 003/117] fix style

---
 cpp/include/cugraph/sampling_functions.hpp    |   5 +-
 cpp/include/cugraph_c/sampling_algorithms.h   |   9 +-
 cpp/src/c_api/neighbor_sampling.cpp           | 131 +++++++++---------
 cpp/src/sampling/neighbor_sampling_impl.hpp   |  43 +++---
 .../sampling/neighbor_sampling_mg_v32_e32.cpp |  10 +-
 .../sampling/neighbor_sampling_mg_v32_e64.cpp |  10 +-
 .../sampling/neighbor_sampling_mg_v64_e64.cpp |  10 +-
 .../sampling/neighbor_sampling_sg_v32_e32.cpp |  10 +-
 .../sampling/neighbor_sampling_sg_v32_e64.cpp |  10 +-
 .../sampling/neighbor_sampling_sg_v64_e64.cpp |  10 +-
 .../dask/sampling/uniform_neighbor_sample.py  |   5 +-
 .../sampling/uniform_neighbor_sample.py       |   8 +-
 .../_cugraph_c/sampling_algorithms.pxd        |   4 +-
 .../pylibcugraph/uniform_neighbor_sample.pyx  |   6 +-
 14 files changed, 149 insertions(+), 122 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 5b884fa611e..7acda3e8d4c 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -130,7 +130,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement                           = true,
@@ -211,7 +212,7 @@ uniform_neighbor_sample(
  * optional int32_t hop, optional label_t label, optional size_t offsets)
  */
 
-// tuple with 3 elements. 1 - edge_type, 2- host span of size_t, 3 - fanout vector 
+// tuple with 3 elements. 1 - edge_type, 2- host span of size_t, 3 - fanout vector
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index e440512c482..6bc9403b2f3 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -324,10 +324,10 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
  */
 // FIXME: internal representation should be tuple instead of pairs - Make it more generic (tuple)
 // cugraph_device_tuple_t, host_device_tuple_t,
-//dictionary, key and array
+// dictionary, key and array
 // translate dictionary to a tuple. Add to the draft PR the PLC layer.
 // Concatenate to build the 3 arrays from the PLC layer
-/// mimic 
+/// mimic
 typedef struct {
   int32_t align_;
 } cugraph_sample_heterogeneous_fanout_t;
@@ -359,8 +359,7 @@ cugraph_error_code_t cugraph_create_heterogeneous_fanout(
  *
  * @param [in]    heterogeneous_fanout The edge type size and fanout values
  */
-void cugraph_heterogeneous_fanout_free(
-  cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
+void cugraph_heterogeneous_fanout_free(cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
 
 /**
  * @brief     Uniform Neighborhood Sampling
@@ -712,7 +711,6 @@ cugraph_error_code_t cugraph_test_uniform_neighborhood_sample_result_create(
  * @return error code
  */
 
-
 cugraph_error_code_t cugraph_select_random_vertices(const cugraph_resource_handle_t* handle,
                                                     const cugraph_graph_t* graph,
                                                     cugraph_rng_state_t* rng_state,
@@ -720,7 +718,6 @@ cugraph_error_code_t cugraph_select_random_vertices(const cugraph_resource_handl
                                                     cugraph_type_erased_device_array_t** vertices,
                                                     cugraph_error_t** error);
 
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 9e6e68871c6..e95f9315b00 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -14,13 +14,14 @@
  * limitations under the License.
  */
 
+#include "c_api/neighbor_sampling.hpp"  // FIXME: Remove this and instead use std::tuple
+
 #include "c_api/abstract_functor.hpp"
 #include "c_api/graph.hpp"
 #include "c_api/properties.hpp"
 #include "c_api/random.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
-#include "c_api/neighbor_sampling.hpp" // FIXME: Remove this and instead use std::tuple
 
 #include <cugraph_c/algorithms.h>
 #include <cugraph_c/sampling_algorithms.h>
@@ -113,7 +114,8 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
       heterogeneous_fan_out_(
-        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)),
+        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(
+          heterogeneous_fan_out)),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       options_(options),
       do_expensive_check_(do_expensive_check)
@@ -221,16 +223,16 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
                 raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
                                                  label_to_comm_rank_->size_}))
             : std::nullopt,
-          (fan_out_ != nullptr)
-            ? std::make_optional<raft::host_span<const int>>(fan_out_->as_type<const int>(), fan_out_->size_)
-            : std::nullopt,
-    
+          (fan_out_ != nullptr) ? std::make_optional<raft::host_span<const int>>(
+                                    fan_out_->as_type<const int>(), fan_out_->size_)
+                                : std::nullopt,
+
           (heterogeneous_fan_out_ != nullptr)
             ? std::make_optional(std::make_tuple(
-              raft::host_span<const int>{heterogeneous_fan_out_->edge_type_id->as_type<int>(),
-                                         heterogeneous_fan_out_->edge_type_id->size_},
-              raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
-                                         heterogeneous_fan_out_->fanout->size_}))
+                raft::host_span<const int>{heterogeneous_fan_out_->edge_type_id->as_type<int>(),
+                                           heterogeneous_fan_out_->edge_type_id->size_},
+                raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
+                                           heterogeneous_fan_out_->fanout->size_}))
             : std::nullopt,
           rng_state_->rng_state_,
           options_.return_hops_,
@@ -769,58 +771,61 @@ struct biased_neighbor_sampling_functor : public cugraph::c_api::abstract_functo
   }
 };
 
-  struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
-    raft::handle_t const& handle_;
-    cugraph::c_api::cugraph_graph_t* graph_;
-    cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
-    cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
-    // FIXME: This type doesn't exist: instead create an 'std::tuple<cugraph_type_erased_host_array_t*>'
-    cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
-
-    create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
-                                ::cugraph_graph_t* graph,
-                                ::cugraph_type_erased_host_array_view_t const* edge_type_size,
-                                ::cugraph_type_erased_host_array_view_t const* fanout)
-      : abstract_functor(),
-        handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
-        graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-        edge_type_size_(
-          reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_size)),
-        fanout_(
-          reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
-    {
-    }
+struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
+  raft::handle_t const& handle_;
+  cugraph::c_api::cugraph_graph_t* graph_;
+  cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
+  cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
+  // FIXME: This type doesn't exist: instead create an
+  // 'std::tuple<cugraph_type_erased_host_array_t*>'
+  cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
+
+  create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
+                                      ::cugraph_graph_t* graph,
+                                      ::cugraph_type_erased_host_array_view_t const* edge_type_size,
+                                      ::cugraph_type_erased_host_array_view_t const* fanout)
+    : abstract_functor(),
+      handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      edge_type_size_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+          edge_type_size)),
+      fanout_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
+  {
+  }
 
-    template <typename vertex_t,
-              typename edge_t,
-              typename weight_t,
-              typename edge_type_type_t,
-              bool store_transposed,
-              bool multi_gpu>
-    void operator()()
-    {
-      // FIXME: Remove this check as it is not necessary
-      if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
-        unsupported();
-      } else {
-        std::vector<int32_t> edge_type_size_copy{(int32_t) edge_type_size_->size_};
-        std::vector<int32_t> fanout_copy{(int32_t) fanout_->size_};
-      
-        raft::copy(
-          edge_type_size_copy.data(), edge_type_size_->as_type<int32_t>(), edge_type_size_->size_, handle_.get_stream());
-      
-        raft::copy(
-          fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
-        
-        // std::tuple (template)
-        // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
-        result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
-          new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, INT32),
-          new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)};
-        
-      }
+  template <typename vertex_t,
+            typename edge_t,
+            typename weight_t,
+            typename edge_type_type_t,
+            bool store_transposed,
+            bool multi_gpu>
+  void operator()()
+  {
+    // FIXME: Remove this check as it is not necessary
+    if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
+      unsupported();
+    } else {
+      std::vector<int32_t> edge_type_size_copy{(int32_t)edge_type_size_->size_};
+      std::vector<int32_t> fanout_copy{(int32_t)fanout_->size_};
+
+      raft::copy(edge_type_size_copy.data(),
+                 edge_type_size_->as_type<int32_t>(),
+                 edge_type_size_->size_,
+                 handle_.get_stream());
+
+      raft::copy(
+        fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
+
+      // std::tuple (template)
+      // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
+      result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
+        new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, INT32),
+        new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)};
     }
-  };
+  }
+};
 }  // namespace
 
 extern "C" cugraph_error_code_t cugraph_sampling_options_create(
@@ -1454,8 +1459,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-
-
 extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
@@ -1464,7 +1467,7 @@ extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
   cugraph_error_t** error)
 {
-  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_size, fanout);  
+  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_size, fanout);
 
   return cugraph::c_api::run_algorithm(graph, functor, heterogeneous_fanout, error);
-}
\ No newline at end of file
+}
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 1989b733eaf..2977084cce2 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -60,8 +60,9 @@ neighbor_sample_impl(
   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,
+  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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -78,20 +79,25 @@ neighbor_sample_impl(
   static_assert(std::is_floating_point_v<bias_t>);
 
   if (fan_out) {
-    CUGRAPH_EXPECTS((*fan_out).size() > 0, "Invalid input argument: number of levels must be non-zero.");
+    CUGRAPH_EXPECTS((*fan_out).size() > 0,
+                    "Invalid input argument: number of levels must be non-zero.");
     CUGRAPH_EXPECTS(
-    (*fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
-    "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
-                                                                              // to store hops
-  } else {    
-    CUGRAPH_EXPECTS(std::accumulate(
-      std::get<0>(*heterogeneous_fan_out).begin(),
-      std::get<0>(*heterogeneous_fan_out).end(), 0) == std::get<1>(*heterogeneous_fan_out).size() && std::get<1>(*heterogeneous_fan_out).size() != 0,
-        "Invalid input argument: number of levels and size must match and should be non zero.");
-    
+      (*fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
+      "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
+                                                                                // to store hops
+  } else {
     CUGRAPH_EXPECTS(
-      std::get<0>(*heterogeneous_fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max())
-        && std::get<1>(*heterogeneous_fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
+      std::accumulate(std::get<0>(*heterogeneous_fan_out).begin(),
+                      std::get<0>(*heterogeneous_fan_out).end(),
+                      0) == std::get<1>(*heterogeneous_fan_out).size() &&
+        std::get<1>(*heterogeneous_fan_out).size() != 0,
+      "Invalid input argument: number of levels and size must match and should be non zero.");
+
+    CUGRAPH_EXPECTS(
+      std::get<0>(*heterogeneous_fan_out).size() <=
+          static_cast<size_t>(std::numeric_limits<int32_t>::max()) &&
+        std::get<1>(*heterogeneous_fan_out).size() <=
+          static_cast<size_t>(std::numeric_limits<int32_t>::max()),
       "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                                 // to store hops
   }
@@ -142,7 +148,9 @@ neighbor_sample_impl(
   level_result_dst_vectors.reserve((*fan_out).size());
   if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve((*fan_out).size()); }
   if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve((*fan_out).size()); }
-  if (level_result_edge_type_vectors) { (*level_result_edge_type_vectors).reserve((*fan_out).size()); }
+  if (level_result_edge_type_vectors) {
+    (*level_result_edge_type_vectors).reserve((*fan_out).size());
+  }
   if (level_result_label_vectors) { (*level_result_label_vectors).reserve((*fan_out).size()); }
 
   rmm::device_uvector<vertex_t> frontier_vertices(0, handle.get_stream());
@@ -382,8 +390,9 @@ uniform_neighbor_sample(
   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,
+  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_e32.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
index 04cba83ac53..a8bc9b015f1 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
index 6ba353ad58c..4bcc93a83d0 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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 a33bff7d3b1..eb99ca0c2a1 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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 44f82c0b0ea..7de1b877b81 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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 69128c02fd2..06ce85e74f3 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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 056730f8100..6237aa1ada0 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
@@ -40,7 +40,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -67,7 +68,8 @@ uniform_neighbor_sample(
   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,
+  std::optional<std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>>
+    heterogeneous_fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement,
@@ -94,7 +96,7 @@ 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,
+  // 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,
@@ -122,7 +124,7 @@ 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,
+  // 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/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index 688a9d973ef..af736ce7846 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -622,7 +622,8 @@ def uniform_neighbor_sample(
         edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
         fanout_vals = (
             numpy.asarray(edge_type_size, dtype="int32"),
-            numpy.asarray(edge_type_fanout_vals, dtype="int32"))
+            numpy.asarray(edge_type_fanout_vals, dtype="int32"),
+        )
     else:
         raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
 
diff --git a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
index 7ba482f36bd..d17e070ad3a 100644
--- a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
@@ -324,9 +324,13 @@ def uniform_neighbor_sample(
         edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
         fanout_vals = (
             numpy.asarray(edge_type_size, dtype="int32"),
-            numpy.asarray(edge_type_fanout_vals, dtype="int32"))
+            numpy.asarray(edge_type_fanout_vals, dtype="int32"),
+        )
     else:
-        raise TypeError("fanout_vals must be a sequence or a dictionary, " f"got: {type(fanout_vals)}")
+        raise TypeError(
+            "fanout_vals must be a sequence or a dictionary, "
+            f"got: {type(fanout_vals)}"
+        )
 
     if "weights" in G.edgelist.edgelist_df:
         weight_t = G.edgelist.edgelist_df["weights"].dtype
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index a9cea82c3d7..c833abae5a9 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -104,10 +104,10 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
             cugraph_type_erased_device_array_t** vertices,
             cugraph_error_t** error
         )
-    
+
     ctypedef struct cugraph_sample_heterogeneous_fanout_t:
         pass
-    
+
     cdef cugraph_error_code_t \
         cugraph_create_heterogeneous_fanout(
             const cugraph_resource_handle_t* handle,
diff --git a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
index fee88b4d16f..c4887414122 100644
--- a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
@@ -238,7 +238,7 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
         assert_AI_type(h_fan_out, "h_fan_out")
         ai_fan_out_ptr = \
             h_fan_out.__array_interface__["data"][0]
-        
+
         fan_out_ptr = \
             cugraph_type_erased_host_array_view_create(
                 <void*>ai_fan_out_ptr,
@@ -257,13 +257,13 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
                 <void*>ai_fan_out_size_ptr,
                 len(h_fan_out[0]),
                 get_c_type_from_numpy_type(h_fan_out[0].dtype))
-    
+
         fan_out_values_ptr = \
             cugraph_type_erased_host_array_view_create(
                 <void*>ai_fan_out_values_ptr,
                 len(h_fan_out[1]),
                 get_c_type_from_numpy_type(h_fan_out[1].dtype))
-        
+
         error_code = cugraph_create_heterogeneous_fanout(
             c_resource_handle_ptr,
             c_graph_ptr,

From f9043508130dcfbe49c716054a1beaf939853a1b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 20 Aug 2024 04:15:13 -0700
Subject: [PATCH 004/117] create one API for both uniform and biased
 neighborhood sampling

---
 cpp/include/cugraph/sampling_functions.hpp    | 124 ++++++++++++++++++
 cpp/include/cugraph_c/sampling_algorithms.h   |  64 +++++++++
 cpp/src/sampling/neighbor_sampling_impl.hpp   |  71 ++++++++++
 .../sampling/neighbor_sampling_mg_v32_e32.cpp |  60 ++++++++-
 .../sampling/neighbor_sampling_mg_v32_e64.cpp |  58 ++++++++
 .../sampling/neighbor_sampling_mg_v64_e64.cpp |  58 ++++++++
 .../sampling/neighbor_sampling_sg_v32_e32.cpp |  58 ++++++++
 .../sampling/neighbor_sampling_sg_v32_e64.cpp |  58 ++++++++
 .../sampling/neighbor_sampling_sg_v64_e64.cpp |  58 ++++++++
 9 files changed, 607 insertions(+), 2 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 7acda3e8d4c..9ee9ce035a3 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -42,6 +42,7 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
 
 /**
  * @brief Uniform Neighborhood Sampling.
+ * # FIXME: Reset all changes to uniform neighborhood sampling
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects from these outgoing neighbors to extract a subgraph.
@@ -350,6 +351,129 @@ biased_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
+
+/**
+ * @brief Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
+ * (optional) will be described below and is dependent upon the input parameters.
+ *
+ * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
+ * label and offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
+ * the label output has values.  This will also result in the output being sorted by vertex label.
+ * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
+ * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
+ * label output has values.  This will also result in the output being sorted by vertex label.  The
+ * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
+ * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
+ * be shuffled so that all data with a particular label will be on the specified rank.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam label_t Type of label. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
+ * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
+ * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
+ * corresponding edge can never be selected.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * for the sampling.
+ * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
+ * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
+ * output rank.  The label span must be sorted in ascending order.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level
+ * @param heterogeneous_fan_out Tuple of host spans defining branching out (fan-out) degree per
+ * source vertex for each level in CSR style format. The first element of the tuple is the offset
+ * array per edge type id and the second element correspond to the fanout values.
+ * @param return_hops boolean flag specifying if the hop information should be returned
+ * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
+ * DEFAULT)
+ * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
+ * multiple times with the same label, it will only be passed once (for each label) as a source
+ * for the next hop.  Default is false.
+ * @param with_replacement boolean flag specifying if random sampling is done with replacement
+ * (true); or, without replacement (false); default = true;
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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,
+  // FIXME: replace 2 optional fan_out type by a std::variant instead
+  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,
+  bool return_hops,
+  bool with_replacement                           = true,
+  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
+  bool dedupe_sources                             = false,
+  bool do_expensive_check                         = false);
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 #if 0
 /* FIXME:
  There are two options to support heterogeneous fanout
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 6bc9403b2f3..fcdb89dee2d 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -476,6 +476,70 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
+/**
+ * @brief     Neighborhood Sampling
+ *
+ * Returns a sample of the neighborhood around specified start vertices with edge biases or not.
+ * Optionally, each start vertex can be associated with a label, allowing the caller to specify
+ * multiple batches of sampling requests in the same function call - which should improve GPU
+ * utilization.
+ *
+ * If label is NULL then all start vertices will be considered part of the same batch and the
+ * return value will not have a label column.
+ *
+ * @param [in]  handle       Handle for accessing resources
+ *  * @param [in,out] rng_state State of the random number generator, updated with each call
+ * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
+ *                           needs to be transposed
+ * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
+ * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
+ * @param [in]  start_vertices Device array of start vertices for the sampling
+ * @param [in]  start_vertex_labels  Device array of start vertex labels for the sampling.  The
+ * labels associated with each start vertex will be included in the output associated with results
+ * that were derived from that start vertex.  We only support label of type INT32. If label is
+ * NULL, the return data will not be labeled.
+ * @param [in]  label_list Device array of the labels included in @p start_vertex_labels.  If
+ * @p label_to_comm_rank is not specified this parameter is ignored.  If specified, label_list
+ * must be sorted in ascending order.
+ * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
+ * particular label should be shuffled in the output.  If not specifed the data is not organized in
+ * output.  If specified then the all data from @p label_list[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_labels is also specified
+ * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
+ * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
+ * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
+ *                           We only support fanout values of type INT32
+ * @param [in]  heterogeneous_fanout  Tuple of host arrays defining the fan out at each step in the
+ * sampling algorithm. in CSR style format. The first element of the tuple is the offset array per
+ * edge type id and the second element correspond to the fanout values.
+ * We only support type INT32 for both the offsets and the fanout values array.
+ * @param [in]  sampling_options
+ *                           Opaque pointer defining the sampling options.
+ * @param [in]  do_expensive_check
+ *                           A flag to run expensive checks for input arguments (if set to true)
+ * @param [out]  result      Output from the uniform_neighbor_sample call
+ * @param [out] error        Pointer to an error object storing details of any error.  Will
+ *                           be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_edge_property_view_t* edge_biases,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error);
+
 /**
  * @deprecated This call should be replaced with cugraph_sample_result_get_majors
  * @brief     Get the source vertices from the sampling algorithm result
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 2977084cce2..c2f6e7bb366 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -476,4 +476,75 @@ biased_neighbor_sample(
     do_expensive_check);
 }
 
+
+
+
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check)
+{
+  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    edge_bias_view,
+    starting_vertices,
+    starting_vertex_labels,
+    label_to_output_comm_rank,
+    fan_out,
+    heterogeneous_fan_out,
+    return_hops,
+    with_replacement,
+    prior_sources_behavior,
+    dedupe_sources,
+    rng_state,
+    do_expensive_check);
+}
+
+
+
+
+
+
+
+
+
+
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
index a8bc9b015f1..1a3e56df26a 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
@@ -96,7 +96,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 +123,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,
@@ -133,4 +131,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
index 4bcc93a83d0..0d3526e202f 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
@@ -133,4 +133,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
index eb99ca0c2a1..a5844fecf80 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
@@ -133,4 +133,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
index 7de1b877b81..c82ff6c2c03 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
@@ -133,4 +133,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
index 06ce85e74f3..ae51c392f5c 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
@@ -133,4 +133,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
index 6237aa1ada0..704313c67fe 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
@@ -133,4 +133,62 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph

From 1fc32c34c90408add19735fdd43d78f083160c1b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 20 Aug 2024 08:56:32 -0700
Subject: [PATCH 005/117] use the same function for both uniform and biased
 nieghborhood sampling

---
 cpp/include/cugraph/sampling_functions.hpp  |   2 +-
 cpp/src/c_api/neighbor_sampling.cpp         | 500 +++++---------------
 cpp/src/sampling/neighbor_sampling_impl.hpp |   6 +-
 3 files changed, 126 insertions(+), 382 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 9ee9ce035a3..d7533a98586 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -423,6 +423,7 @@ biased_neighbor_sample(
  * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
  * optional int32_t hop, optional label_t label, optional size_t offsets)
  */
+// FIXME: Add flag for bias=True/False
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
@@ -451,7 +452,6 @@ 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,
-  // FIXME: replace 2 optional fan_out type by a std::variant instead
   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,
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index e95f9315b00..212f4736d42 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -66,9 +66,11 @@ struct cugraph_sample_result_t {
 
 namespace {
 
-struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
+struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
+  cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
@@ -76,14 +78,15 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out_{nullptr};
-  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
-  uniform_neighbor_sampling_functor(
+  neighbor_sampling_functor(
     cugraph_resource_handle_t const* handle,
+    cugraph_rng_state_t* rng_state,
     cugraph_graph_t* graph,
+    cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
     cugraph_type_erased_device_array_view_t const* start_vertex_labels,
     cugraph_type_erased_device_array_view_t const* label_list,
@@ -91,12 +94,14 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
     cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
     cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out,
-    cugraph_rng_state_t* rng_state,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      edge_biases_(
+        reinterpret_cast<cugraph::c_api::cugraph_edge_property_view_t const*>(edge_biases)),
       start_vertices_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           start_vertices)),
@@ -116,7 +121,6 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
       heterogeneous_fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(
           heterogeneous_fan_out)),
-      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       options_(options),
       do_expensive_check_(do_expensive_check)
   {
@@ -163,6 +167,11 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
 
       auto number_map = reinterpret_cast<rmm::device_uvector<vertex_t>*>(graph_->number_map_);
 
+      auto edge_biases =
+        edge_biases_ ? reinterpret_cast<cugraph::edge_property_view_t<edge_t, weight_t const*>*>(
+                         edge_biases_->edge_property_)
+                     : nullptr;
+
       rmm::device_uvector<vertex_t> start_vertices(start_vertices_->size_, handle_.get_stream());
       raft::copy(start_vertices.data(),
                  start_vertices_->as_type<vertex_t>(),
@@ -205,12 +214,16 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
         do_expensive_check_);
 
       auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-        cugraph::uniform_neighbor_sample(
+        cugraph::neighbor_sample(
           handle_,
+          rng_state_->rng_state_,
           graph_view,
           (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
           (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
           (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+          // FIXME: add bias flag to differentiate between bias and uniform neighbor sample
+          // bias ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt ***
+          (edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view()),
           raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
           (start_vertex_labels_ != nullptr)
             ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
@@ -234,7 +247,6 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
                 raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
                                            heterogeneous_fan_out_->fanout->size_}))
             : std::nullopt,
-          rng_state_->rng_state_,
           options_.return_hops_,
           options_.with_replacement_,
           options_.prior_sources_behavior_,
@@ -421,355 +433,6 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   }
 };
 
-struct biased_neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
-  raft::handle_t const& handle_;
-  cugraph::c_api::cugraph_graph_t* graph_{nullptr};
-  cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
-  cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
-  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
-  cugraph::c_api::cugraph_sampling_options_t options_{};
-  bool do_expensive_check_{false};
-  cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
-
-  biased_neighbor_sampling_functor(
-    cugraph_resource_handle_t const* handle,
-    cugraph_graph_t* graph,
-    cugraph_edge_property_view_t const* edge_biases,
-    cugraph_type_erased_device_array_view_t const* start_vertices,
-    cugraph_type_erased_device_array_view_t const* start_vertex_labels,
-    cugraph_type_erased_device_array_view_t const* label_list,
-    cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
-    cugraph_type_erased_device_array_view_t const* label_offsets,
-    cugraph_type_erased_host_array_view_t const* fan_out,
-    cugraph_rng_state_t* rng_state,
-    cugraph::c_api::cugraph_sampling_options_t options,
-    bool do_expensive_check)
-    : abstract_functor(),
-      handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
-      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-      edge_biases_(
-        reinterpret_cast<cugraph::c_api::cugraph_edge_property_view_t const*>(edge_biases)),
-      start_vertices_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          start_vertices)),
-      start_vertex_labels_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          start_vertex_labels)),
-      label_list_(reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-        label_list)),
-      label_to_comm_rank_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          label_to_comm_rank)),
-      label_offsets_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          label_offsets)),
-      fan_out_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
-      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
-      options_(options),
-      do_expensive_check_(do_expensive_check)
-  {
-  }
-
-  template <typename vertex_t,
-            typename edge_t,
-            typename weight_t,
-            typename edge_type_t,
-            bool store_transposed,
-            bool multi_gpu>
-  void operator()()
-  {
-    using label_t = int32_t;
-
-    // FIXME: Think about how to handle SG vice MG
-    if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
-      unsupported();
-    } else {
-      // uniform_nbr_sample expects store_transposed == false
-      if constexpr (store_transposed) {
-        error_code_ = cugraph::c_api::
-          transpose_storage<vertex_t, edge_t, weight_t, store_transposed, multi_gpu>(
-            handle_, graph_, error_.get());
-        if (error_code_ != CUGRAPH_SUCCESS) return;
-      }
-
-      auto graph =
-        reinterpret_cast<cugraph::graph_t<vertex_t, edge_t, false, multi_gpu>*>(graph_->graph_);
-
-      auto graph_view = graph->view();
-
-      auto edge_weights = reinterpret_cast<
-        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
-                                 weight_t>*>(graph_->edge_weights_);
-
-      auto edge_ids = reinterpret_cast<
-        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
-                                 edge_t>*>(graph_->edge_ids_);
-
-      auto edge_types = reinterpret_cast<
-        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
-                                 edge_type_t>*>(graph_->edge_types_);
-
-      auto number_map = reinterpret_cast<rmm::device_uvector<vertex_t>*>(graph_->number_map_);
-
-      auto edge_biases =
-        edge_biases_ ? reinterpret_cast<cugraph::edge_property_view_t<edge_t, weight_t const*>*>(
-                         edge_biases_->edge_property_)
-                     : nullptr;
-
-      rmm::device_uvector<vertex_t> start_vertices(start_vertices_->size_, handle_.get_stream());
-      raft::copy(start_vertices.data(),
-                 start_vertices_->as_type<vertex_t>(),
-                 start_vertices.size(),
-                 handle_.get_stream());
-
-      std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
-
-      if (start_vertex_labels_ != nullptr) {
-        start_vertex_labels =
-          rmm::device_uvector<label_t>{start_vertex_labels_->size_, handle_.get_stream()};
-        raft::copy(start_vertex_labels->data(),
-                   start_vertex_labels_->as_type<label_t>(),
-                   start_vertex_labels_->size_,
-                   handle_.get_stream());
-      }
-
-      if constexpr (multi_gpu) {
-        if (start_vertex_labels) {
-          std::tie(start_vertices, *start_vertex_labels) =
-            cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
-              handle_, std::move(start_vertices), std::move(*start_vertex_labels));
-        } else {
-          start_vertices =
-            cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
-              handle_, std::move(start_vertices));
-        }
-      }
-
-      //
-      // Need to renumber start_vertices
-      //
-      cugraph::renumber_local_ext_vertices<vertex_t, multi_gpu>(
-        handle_,
-        start_vertices.data(),
-        start_vertices.size(),
-        number_map->data(),
-        graph_view.local_vertex_partition_range_first(),
-        graph_view.local_vertex_partition_range_last(),
-        do_expensive_check_);
-
-      auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-        cugraph::biased_neighbor_sample(
-          handle_,
-          graph_view,
-          (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-          (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-          (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-          (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
-          raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-          (start_vertex_labels_ != nullptr)
-            ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                   start_vertex_labels->size())
-            : std::nullopt,
-          (label_list_ != nullptr)
-            ? std::make_optional(std::make_tuple(
-                raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                 label_list_->size_},
-                raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                 label_to_comm_rank_->size_}))
-            : std::nullopt,
-          raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
-          rng_state_->rng_state_,
-          options_.return_hops_,
-          options_.with_replacement_,
-          options_.prior_sources_behavior_,
-          options_.dedupe_sources_,
-          do_expensive_check_);
-
-      std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
-
-      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
-                                                            src.data(),
-                                                            src.size(),
-                                                            number_map->data(),
-                                                            vertex_partition_lasts,
-                                                            do_expensive_check_);
-
-      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
-                                                            dst.data(),
-                                                            dst.size(),
-                                                            number_map->data(),
-                                                            vertex_partition_lasts,
-                                                            do_expensive_check_);
-
-      std::optional<rmm::device_uvector<vertex_t>> majors{std::nullopt};
-      rmm::device_uvector<vertex_t> minors(0, handle_.get_stream());
-      std::optional<rmm::device_uvector<size_t>> major_offsets{std::nullopt};
-
-      std::optional<rmm::device_uvector<size_t>> label_hop_offsets{std::nullopt};
-
-      std::optional<rmm::device_uvector<vertex_t>> renumber_map{std::nullopt};
-      std::optional<rmm::device_uvector<size_t>> renumber_map_offsets{std::nullopt};
-
-      bool src_is_major = (options_.compression_type_ == cugraph_compression_type_t::CSR) ||
-                          (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
-                          (options_.compression_type_ == cugraph_compression_type_t::COO);
-
-      if (options_.renumber_results_) {
-        if (options_.compression_type_ == cugraph_compression_type_t::COO) {
-          // COO
-
-          rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
-          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-          std::tie(output_majors,
-                   minors,
-                   wgt,
-                   edge_id,
-                   edge_type,
-                   label_hop_offsets,
-                   output_renumber_map,
-                   renumber_map_offsets) =
-            cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
-              handle_,
-              std::move(src),
-              std::move(dst),
-              std::move(wgt),
-              std::move(edge_id),
-              std::move(edge_type),
-              std::move(hop),
-              options_.retain_seeds_
-                ? std::make_optional(raft::device_span<vertex_t const>{
-                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
-                : std::nullopt,
-              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
-                                     : std::nullopt,
-              offsets ? std::make_optional(
-                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
-                      : std::nullopt,
-              edge_label ? edge_label->size() : size_t{1},
-              hop ? fan_out_->size_ : size_t{1},
-              src_is_major,
-              do_expensive_check_);
-
-          majors.emplace(std::move(output_majors));
-          renumber_map.emplace(std::move(output_renumber_map));
-        } else {
-          // (D)CSC, (D)CSR
-
-          bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
-                                 (options_.compression_type_ == cugraph_compression_type_t::DCSC);
-
-          rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
-          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-          std::tie(majors,
-                   output_major_offsets,
-                   minors,
-                   wgt,
-                   edge_id,
-                   edge_type,
-                   label_hop_offsets,
-                   output_renumber_map,
-                   renumber_map_offsets) =
-            cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
-              handle_,
-              std::move(src),
-              std::move(dst),
-              std::move(wgt),
-              std::move(edge_id),
-              std::move(edge_type),
-              std::move(hop),
-              options_.retain_seeds_
-                ? std::make_optional(raft::device_span<vertex_t const>{
-                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
-                : std::nullopt,
-              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
-                                     : std::nullopt,
-              offsets ? std::make_optional(
-                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
-                      : std::nullopt,
-              edge_label ? edge_label->size() : size_t{1},
-              hop ? fan_out_->size_ : size_t{1},
-              src_is_major,
-              options_.compress_per_hop_,
-              doubly_compress,
-              do_expensive_check_);
-
-          renumber_map.emplace(std::move(output_renumber_map));
-          major_offsets.emplace(std::move(output_major_offsets));
-        }
-
-        // These are now represented by label_hop_offsets
-        hop.reset();
-        offsets.reset();
-      } else {
-        if (options_.compression_type_ != cugraph_compression_type_t::COO) {
-          CUGRAPH_FAIL("Can only use COO format if not renumbering");
-        }
-
-        std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
-          cugraph::sort_sampled_edgelist(handle_,
-                                         std::move(src),
-                                         std::move(dst),
-                                         std::move(wgt),
-                                         std::move(edge_id),
-                                         std::move(edge_type),
-                                         std::move(hop),
-                                         offsets
-                                           ? std::make_optional(raft::device_span<size_t const>{
-                                               offsets->data(), offsets->size()})
-                                           : std::nullopt,
-                                         edge_label ? edge_label->size() : size_t{1},
-                                         hop ? fan_out_->size_ : size_t{1},
-                                         src_is_major,
-                                         do_expensive_check_);
-
-        majors.emplace(std::move(src));
-        minors = std::move(dst);
-
-        hop.reset();
-        offsets.reset();
-      }
-
-      result_ = new cugraph::c_api::cugraph_sample_result_t{
-        (major_offsets)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*major_offsets, SIZE_T)
-          : nullptr,
-        (majors)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*majors, graph_->vertex_type_)
-          : nullptr,
-        new cugraph::c_api::cugraph_type_erased_device_array_t(minors, graph_->vertex_type_),
-        (edge_id)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*edge_id, graph_->edge_type_)
-          : nullptr,
-        (edge_type) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
-                        *edge_type, graph_->edge_type_id_type_)
-                    : nullptr,
-        (wgt) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*wgt, graph_->weight_type_)
-              : nullptr,
-        (hop) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*hop, INT32)
-              : nullptr,  // FIXME get rid of this
-        (label_hop_offsets)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*label_hop_offsets, SIZE_T)
-          : nullptr,
-        (edge_label)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32)
-          : nullptr,
-        (renumber_map) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
-                           renumber_map.value(), graph_->vertex_type_)
-                       : nullptr,
-        (renumber_map_offsets) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
-                                   renumber_map_offsets.value(), SIZE_T)
-                               : nullptr};
-    }
-  }
-};
 
 struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
@@ -1366,18 +1029,93 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
-  uniform_neighbor_sampling_functor functor{handle,
-                                            graph,
-                                            start_vertices,
-                                            start_vertex_labels,
-                                            label_list,
-                                            label_to_comm_rank,
-                                            label_offsets,
-                                            fan_out,
-                                            heterogeneous_fanout,
-                                            rng_state,
-                                            std::move(options_cpp),
-                                            do_expensive_check};
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    nullptr,
+                                    start_vertices,
+                                    start_vertex_labels,
+                                    label_list,
+                                    label_to_comm_rank,
+                                    label_offsets,
+                                    fan_out,
+                                    heterogeneous_fanout,
+                                    std::move(options_cpp),
+                                    do_expensive_check};
+  return cugraph::c_api::run_algorithm(graph, functor, result, error);
+}
+
+// FIXME: Add flag for bias=True/False
+cugraph_error_code_t cugraph_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_edge_property_view_t* edge_biases,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error)
+{
+  auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
+
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "must specify label_offsets if retain_seeds is true",
+               *error);
+
+  CAPI_EXPECTS((start_vertex_labels == nullptr) ||
+                 (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                    start_vertex_labels)
+                    ->type_ == INT32),
+               CUGRAPH_INVALID_INPUT,
+               "start_vertex_labels should be of type int",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_labels != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless start_vertex_labels is also specified",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless label_list is also specified",
+               *error);
+
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
+                 reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                   start_vertices)
+                   ->type_,
+               CUGRAPH_INVALID_INPUT,
+               "vertex type of graph and start_vertices must match",
+               *error);
+
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out should be of type int",
+    *error);
+
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    edge_biases,
+                                    start_vertices,
+                                    start_vertex_labels,
+                                    label_list,
+                                    label_to_comm_rank,
+                                    label_offsets,
+                                    fan_out,
+                                    heterogeneous_fanout,
+                                    std::move(options_cpp),
+                                    do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
@@ -1399,6 +1137,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
+  // FIXME: Check if bias = True before checking the statement below
   CAPI_EXPECTS(
     (edge_biases != nullptr) ||
       (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
@@ -1444,18 +1183,19 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
-  biased_neighbor_sampling_functor functor{handle,
-                                           graph,
-                                           edge_biases,
-                                           start_vertices,
-                                           start_vertex_labels,
-                                           label_list,
-                                           label_to_comm_rank,
-                                           label_offsets,
-                                           fan_out,
-                                           rng_state,
-                                           std::move(options_cpp),
-                                           do_expensive_check};
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    edge_biases,
+                                    start_vertices,
+                                    start_vertex_labels,
+                                    label_list,
+                                    label_to_comm_rank,
+                                    label_offsets,
+                                    fan_out,
+                                    nullptr,
+                                    std::move(options_cpp),
+                                    do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index c2f6e7bb366..b82e306cc4b 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -62,7 +62,6 @@ neighbor_sample_impl(
   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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -86,12 +85,17 @@ neighbor_sample_impl(
       "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                                 // to store hops
   } else {
+    /*
+    FIXME: Since we are dealing with offsets and size, just check that the last element of the 
+    first host array == size of the second array.
+
     CUGRAPH_EXPECTS(
       std::accumulate(std::get<0>(*heterogeneous_fan_out).begin(),
                       std::get<0>(*heterogeneous_fan_out).end(),
                       0) == std::get<1>(*heterogeneous_fan_out).size() &&
         std::get<1>(*heterogeneous_fan_out).size() != 0,
       "Invalid input argument: number of levels and size must match and should be non zero.");
+    */
 
     CUGRAPH_EXPECTS(
       std::get<0>(*heterogeneous_fan_out).size() <=

From 8fc21f8f219aab89656db54ba4330b24b6ef08ea Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 20 Aug 2024 08:57:13 -0700
Subject: [PATCH 006/117] add support for heterogenous fanout support at the
 plc layer and consolidate neighborhood sampling functions

---
 .../pylibcugraph/pylibcugraph/CMakeLists.txt  |   1 +
 python/pylibcugraph/pylibcugraph/__init__.py  |   2 +
 .../_cugraph_c/sampling_algorithms.pxd        |  18 +
 .../pylibcugraph/neighbor_sample.pyx          | 514 ++++++++++++++++++
 4 files changed, 535 insertions(+)
 create mode 100644 python/pylibcugraph/pylibcugraph/neighbor_sample.pyx

diff --git a/python/pylibcugraph/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
index 90fce23282e..7b28ff40059 100644
--- a/python/pylibcugraph/pylibcugraph/CMakeLists.txt
+++ b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
@@ -54,6 +54,7 @@ set(cython_sources
     triangle_count.pyx
     two_hop_neighbors.pyx
     uniform_neighbor_sample.pyx
+    neighbor_sample.pyx
     uniform_random_walks.pyx
     utils.pyx
     weakly_connected_components.pyx
diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py
index b67acc8bbfc..7db0a218dab 100644
--- a/python/pylibcugraph/pylibcugraph/__init__.py
+++ b/python/pylibcugraph/pylibcugraph/__init__.py
@@ -39,6 +39,8 @@
 
 from pylibcugraph.uniform_neighbor_sample import uniform_neighbor_sample
 
+from pylibcugraph.uniform_neighbor_sample import neighbor_sample
+
 from pylibcugraph.core_number import core_number
 
 from pylibcugraph.k_core import k_core
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index c833abae5a9..c47621db191 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -64,6 +64,24 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_error_t** error
     )
 
+    cdef cugraph_error_code_t cugraph_neighbor_sample(
+        const cugraph_resource_handle_t* handle,
+        cugraph_rng_state_t* rng_state,
+        cugraph_graph_t* graph,
+        const cugraph_edge_property_view_t* edge_biases,
+        const cugraph_type_erased_device_array_view_t* start_vertices,
+        const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+        const cugraph_type_erased_device_array_view_t* label_list,
+        const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+        const cugraph_type_erased_device_array_view_t* label_offsets,
+        const cugraph_type_erased_host_array_view_t* fan_out,
+        const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+        const cugraph_sampling_options_t* options,
+        bool_t do_expensive_check,
+        cugraph_sample_result_t** result,
+        cugraph_error_t** error
+    )
+
     cdef cugraph_error_code_t cugraph_biased_neighbor_sample(
         const cugraph_resource_handle_t* handle,
         cugraph_graph_t* graph,
diff --git a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
new file mode 100644
index 00000000000..d3f8bbae9da
--- /dev/null
+++ b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
@@ -0,0 +1,514 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    bool_t,
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.properties cimport (
+    cugraph_edge_property_view_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.algorithms cimport (
+    cugraph_sample_result_t,
+    cugraph_prior_sources_behavior_t,
+    cugraph_compression_type_t,
+    cugraph_sampling_options_t,
+    cugraph_sampling_options_create,
+    cugraph_sampling_options_free,
+    cugraph_sampling_set_with_replacement,
+    cugraph_sampling_set_return_hops,
+    cugraph_sampling_set_prior_sources_behavior,
+    cugraph_sampling_set_dedupe_sources,
+    cugraph_sampling_set_renumber_results,
+    cugraph_sampling_set_compress_per_hop,
+    cugraph_sampling_set_compression_type,
+    cugraph_sampling_set_retain_seeds,
+)
+from pylibcugraph._cugraph_c.sampling_algorithms cimport (
+    cugraph_neighbor_sample,
+    cugraph_sample_heterogeneous_fanout_t,
+    cugraph_create_heterogeneous_fanout,
+    cugraph_heterogeneous_fanout_free,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+)
+from pylibcugraph.internal_types.sampling_result cimport (
+    SamplingResult,
+)
+from pylibcugraph._cugraph_c.random cimport (
+    cugraph_rng_state_t
+)
+from pylibcugraph.random cimport (
+    CuGraphRandomState
+)
+import warnings
+
+# TODO accept cupy/numpy random state in addition to raw seed.
+def neighbor_sample(ResourceHandle resource_handle,
+                    _GPUGraph input_graph,
+                    start_list,
+                    h_fan_out,
+                    bool_t with_replacement,
+                    bool_t do_expensive_check,
+                    with_edge_properties=False,
+                    batch_id_list=None,
+                    label_list=None,
+                    label_to_output_comm_rank=None,
+                    label_offsets=None,
+                    prior_sources_behavior=None,
+                    deduplicate_sources=False,
+                    return_hops=False,
+                    renumber=False,
+                    retain_seeds=False,
+                    compression='COO',
+                    compress_per_hop=False,
+                    random_state=None,
+                    return_dict=False,):
+    """
+    # FIXME: Deprecate uniform_neighbor_sample
+    Does both uniform and biased neighborhood sampling, which samples nodes from
+    a graph based on the current node's neighbors, with a corresponding fanout
+    value at each hop.
+
+    # FIXME: biased neighbor sampling is not yet exposed to PLC.
+
+    Parameters
+    ----------
+    resource_handle: ResourceHandle
+        Handle to the underlying device and host resources needed for
+        referencing data and running algorithms.
+
+    input_graph : SGGraph or MGGraph
+        The input graph, for either Single or Multi-GPU operations.
+
+    start_list: device array type
+        Device array containing the list of starting vertices for sampling.
+
+    edge_biases: FIXE: update this - Create edge_biases of type
+    'cugraph_edge_property_view_t'
+
+    h_fan_out: numpy array type or tuple of numpy array type
+        Device array containing the brancing out (fan-out) degrees per
+        starting vertex for each hop level.
+
+    with_replacement: bool
+        If true, sampling procedure is done with replacement (the same vertex
+        can be selected multiple times in the same step).
+
+    do_expensive_check: bool
+        If True, performs more extensive tests on the inputs to ensure
+        validitity, at the expense of increased run time.
+
+    with_edge_properties: bool
+        If True, returns the edge properties of each edges along with the
+        edges themselves.  Will result in an error if the provided graph
+        does not have edge properties.
+
+    batch_id_list: list[int32] (Optional)
+        List of int32 batch ids that is returned with each edge.  Optional
+        argument, defaults to NULL, returning nothing.
+
+    label_list: list[int32] (Optional)
+        List of unique int32 batch ids.  Required if also passing the
+        label_to_output_comm_rank flag.  Default to NULL (does nothing)
+
+    label_to_output_comm_rank: list[int32] (Optional)
+        Maps the unique batch ids in label_list to the rank of the
+        worker that should hold results for that batch id.
+        Defaults to NULL (does nothing)
+
+    label_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    prior_sources_behavior: str (Optional)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+
+    deduplicate_sources: bool (Optional)
+        If True, will deduplicate the source list before sampling.
+        Defaults to False.
+
+    renumber: bool (Optional)
+        If True, will renumber the sources and destinations on a
+        per-batch basis and return the renumber map and batch offsets
+        in additional to the standard returns.
+
+    retain_seeds: bool (Optional)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+        Defaults to False.
+
+    compression: str (Optional)
+        Options: COO (default), CSR, CSC, DCSR, DCSR
+        Sets the compression format for the returned samples.
+
+    compress_per_hop: bool (Optional)
+        If False (default), will create a compressed edgelist for the
+        entire batch.
+        If True, will create a separate compressed edgelist per hop within
+        a batch.
+
+    random_state: int (Optional)
+        Random state to use when generating samples.  Optional argument,
+        defaults to a hash of process id, time, and hostname.
+        (See pylibcugraph.random.CuGraphRandomState)
+
+    return_dict: bool (Optional)
+        Whether to return a dictionary instead of a tuple.
+        Optional argument, defaults to False, returning a tuple.
+        This argument will eventually be deprecated in favor
+        of always returning a dictionary.
+
+    Returns
+    -------
+    A tuple of device arrays, where the first and second items in the tuple
+    are device arrays containing the starting and ending vertices of each
+    walk respectively, the third item in the tuple is a device array
+    containing the start labels, and the fourth item in the tuple is a device
+    array containing the indices for reconstructing paths.
+
+    If renumber was set to True, then the fifth item in the tuple is a device
+    array containing the renumber map, and the sixth item in the tuple is a
+    device array containing the renumber map offsets (which delineate where
+    the renumber map for each batch starts).
+
+    """
+    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
+        resource_handle.c_resource_handle_ptr
+    )
+
+    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fanout_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+
+    cdef bool_t c_deduplicate_sources = deduplicate_sources
+    cdef bool_t c_return_hops = return_hops
+    cdef bool_t c_renumber = renumber
+    cdef bool_t c_compress_per_hop = compress_per_hop
+
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+    cdef uintptr_t ai_fan_out_size_ptr
+    cdef uintptr_t ai_fan_out_values_ptr
+
+
+    assert_CAI_type(start_list, "start_list")
+    assert_CAI_type(batch_id_list, "batch_id_list", True)
+    assert_CAI_type(label_list, "label_list", True)
+    assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
+    assert_CAI_type(label_offsets, "label_offsets", True)
+    if not isinstance(h_fan_out, tuple):
+        assert_AI_type(h_fan_out, "h_fan_out")
+        ai_fan_out_ptr = \
+            h_fan_out.__array_interface__["data"][0]
+
+        fan_out_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_ptr,
+                len(h_fan_out),
+                get_c_type_from_numpy_type(h_fan_out.dtype))
+    else:
+        assert_AI_type(h_fan_out[0], "h_fan_out_size")
+        assert_AI_type(h_fan_out[1], "h_fan_out_values")
+        ai_fan_out_size_ptr = \
+            h_fan_out[0].__array_interface__["data"][0]
+        ai_fan_out_values_ptr = \
+            h_fan_out[1].__array_interface__["data"][0]
+
+        fan_out_size_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_size_ptr,
+                len(h_fan_out[0]),
+                get_c_type_from_numpy_type(h_fan_out[0].dtype))
+
+        fan_out_values_ptr = \
+            cugraph_type_erased_host_array_view_create(
+                <void*>ai_fan_out_values_ptr,
+                len(h_fan_out[1]),
+                get_c_type_from_numpy_type(h_fan_out[1].dtype))
+
+        error_code = cugraph_create_heterogeneous_fanout(
+            c_resource_handle_ptr,
+            c_graph_ptr,
+            fan_out_size_ptr,
+            fan_out_values_ptr,
+            &heterogeneous_fanout_ptr,
+            &error_ptr
+        )
+
+    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fanout")
+
+    cdef cugraph_sample_result_t* result_ptr
+
+    cdef uintptr_t cai_start_ptr = \
+        start_list.__cuda_array_interface__["data"][0]
+
+    cdef uintptr_t cai_batch_id_ptr
+    if batch_id_list is not None:
+        cai_batch_id_ptr = \
+            batch_id_list.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_list_ptr
+    if label_list is not None:
+        cai_label_list_ptr = \
+            label_list.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_to_output_comm_rank_ptr
+    if label_to_output_comm_rank is not None:
+        cai_label_to_output_comm_rank_ptr = \
+            label_to_output_comm_rank.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_offsets_ptr
+    if label_offsets is not None:
+        cai_label_offsets_ptr = \
+            label_offsets.__cuda_array_interface__['data'][0]
+
+    cdef cugraph_type_erased_device_array_view_t* start_ptr = \
+        cugraph_type_erased_device_array_view_create(
+            <void*>cai_start_ptr,
+            len(start_list),
+            get_c_type_from_numpy_type(start_list.dtype))
+
+    cdef cugraph_type_erased_device_array_view_t* batch_id_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if batch_id_list is not None:
+        batch_id_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_batch_id_ptr,
+                len(batch_id_list),
+                get_c_type_from_numpy_type(batch_id_list.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_list_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if label_list is not None:
+        label_list_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_list_ptr,
+                len(label_list),
+                get_c_type_from_numpy_type(label_list.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_to_output_comm_rank_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if label_to_output_comm_rank is not None:
+        label_to_output_comm_rank_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_to_output_comm_rank_ptr,
+                len(label_to_output_comm_rank),
+                get_c_type_from_numpy_type(label_to_output_comm_rank.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if retain_seeds:
+        if label_offsets is None:
+            raise ValueError("Must provide label offsets if retain_seeds is True")
+        label_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_offsets_ptr,
+                len(label_offsets),
+                get_c_type_from_numpy_type(label_offsets.dtype)
+            )
+
+    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
+
+    cdef cugraph_rng_state_t* rng_state_ptr = \
+        cg_rng_state.rng_state_ptr
+
+    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
+    if prior_sources_behavior is None:
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
+    elif prior_sources_behavior == 'carryover':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
+    elif prior_sources_behavior == 'exclude':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
+    else:
+        raise ValueError(
+            f'Invalid option {prior_sources_behavior}'
+            ' for prior sources behavior'
+        )
+
+    cdef cugraph_compression_type_t compression_behavior_e
+    if compression is None or compression == 'COO':
+        compression_behavior_e = cugraph_compression_type_t.COO
+    elif compression == 'CSR':
+        compression_behavior_e = cugraph_compression_type_t.CSR
+    elif compression == 'CSC':
+        compression_behavior_e = cugraph_compression_type_t.CSC
+    elif compression == 'DCSR':
+        compression_behavior_e = cugraph_compression_type_t.DCSR
+    elif compression == 'DCSC':
+        compression_behavior_e = cugraph_compression_type_t.DCSC
+    else:
+        raise ValueError(
+            f'Invalid option {compression}'
+            ' for compression type'
+        )
+
+    cdef cugraph_sampling_options_t* sampling_options
+    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
+
+    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
+    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
+    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
+    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
+    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
+    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
+    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
+    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
+
+    error_code = cugraph_neighbor_sample(
+        c_resource_handle_ptr,
+        rng_state_ptr,
+        c_graph_ptr,
+        <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
+        start_ptr,
+        batch_id_ptr,
+        label_list_ptr,
+        label_to_output_comm_rank_ptr,
+        label_offsets_ptr,
+        fan_out_ptr,
+        heterogeneous_fanout_ptr,
+        sampling_options,
+        do_expensive_check,
+        &result_ptr,
+        &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_neighbor_sample")
+
+    # Free the sampling options
+    cugraph_sampling_options_free(sampling_options)
+
+    if isinstance(h_fan_out, tuple):
+        cugraph_heterogeneous_fanout_free(heterogeneous_fanout_ptr)
+
+    # Free the two input arrays that are no longer needed.
+    cugraph_type_erased_device_array_view_free(start_ptr)
+    cugraph_type_erased_host_array_view_free(fan_out_ptr)
+    if batch_id_list is not None:
+        cugraph_type_erased_device_array_view_free(batch_id_ptr)
+    if label_offsets is not None:
+        cugraph_type_erased_device_array_view_free(label_offsets_ptr)
+
+    # Have the SamplingResult instance assume ownership of the result data.
+    result = SamplingResult()
+    result.set_ptr(result_ptr)
+
+    # Get cupy "views" of the individual arrays to return. These each increment
+    # the refcount on the SamplingResult instance which will keep the data alive
+    # until all references are removed and the GC runs.
+    # TODO Return everything that isn't null in release 23.12
+    if with_edge_properties:
+        cupy_majors = result.get_majors()
+        cupy_major_offsets = result.get_major_offsets()
+        cupy_minors = result.get_minors()
+        cupy_edge_weights = result.get_edge_weights()
+        cupy_edge_ids = result.get_edge_ids()
+        cupy_edge_types = result.get_edge_types()
+        cupy_batch_ids = result.get_batch_ids()
+        cupy_label_hop_offsets = result.get_label_hop_offsets()
+
+        if renumber:
+            cupy_renumber_map = result.get_renumber_map()
+            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+            # TODO drop the placeholder for hop ids in release 23.12
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': None,
+                    'renumber_map': cupy_renumber_map,
+                    'renumber_map_offsets': cupy_renumber_map_offsets
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+        else:
+            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': cupy_hop_ids,
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
+
+    else:
+        # TODO this is deprecated, remove it in release 23.12
+        warnings.warn(
+            "Calling uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
+            "with_edge_properties=True option, returning whatever properties are in the graph.",
+            FutureWarning,
+        )
+
+        cupy_sources = result.get_sources()
+        cupy_destinations = result.get_destinations()
+        cupy_indices = result.get_indices()
+
+        if return_dict:
+            return {
+                'sources': cupy_sources,
+                'destinations': cupy_destinations,
+                'indices': cupy_indices
+            }
+        else:
+            return (cupy_sources, cupy_destinations, cupy_indices)

From 01a57f39e19f2a434a0a07b2c26ae4881b7d43de Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 20 Aug 2024 09:06:26 -0700
Subject: [PATCH 007/117] remove outdated codes

---
 cpp/include/cugraph/sampling_functions.hpp | 233 ---------------------
 1 file changed, 233 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index d7533a98586..aee0c7f240e 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -140,114 +140,6 @@ uniform_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
-#if 0
-/* FIXME:
- There are two options to support heterogeneous fanout
- 1) Create a new function 'heterogeneous_uniform_neighbor_sample' which will take as input
-      only heterogeneous fanout type. Drawback: code redundancy
- 2) Update 'uniform_neighbor_sample' to support both fanout types
-*/
-/**
- * @brief Heterogeneous uniform Neighborhood Sampling.
- *
- * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects from these outgoing neighbors to extract a subgraph.
- *
- * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
- * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
- * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
- * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
- * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
- * (optional) will be described below and is dependent upon the input parameters.
- *
- * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
- * label and offsets values in the return set will be std::nullopt.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
- * the label output has values.  This will also result in the output being sorted by vertex label.
- * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
- * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
- * label output has values.  This will also result in the output being sorted by vertex label.  The
- * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
- * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
- * be shuffled so that all data with a particular label will be on the specified rank.
- *
- * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
- * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
- * @tparam weight_t Type of edge weights. Needs to be a floating point type.
- * @tparam edge_type_t Type of edge type. Needs to be an integral type.
- * @tparam label_t Type of label. Needs to be an integral type.
- * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
- * true) are major indices
- * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
- * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
- * handles to various CUDA libraries) to run graph algorithms.
- * @param graph_view Graph View object to generate NBR Sampling on.
- * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
- * @param edge_id_view Optional view object holding edge ids for @p graph_view.
- * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param starting_vertices Device span of starting vertex IDs for the sampling.
- * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
- * for the sampling.
- * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
- * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
- * output rank.  The label span must be sorted in ascending order.
- * @param fan_out Tuple of host spans mapping each edge type to fanout values. Element 0
- * of the tuple defines the size of the fanout per edge type while element 1 defines the branching out
- * (fan-out) degree per edge type for each level.
- * @param rng_state A pre-initialized raft::RngState object for generating random numbers
- * @param return_hops boolean flag specifying if the hop information should be returned
- * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
- * DEFAULT)
- * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
- * multiple times with the same label, it will only be passed once (for each label) as a source
- * for the next hop.  Default is false.
- * @param with_replacement boolean flag specifying if random sampling is done with replacement
- * (true); or, without replacement (false); default = true;
- * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
- * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
- * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
- * optional int32_t hop, optional label_t label, optional size_t offsets)
- */
-
-// tuple with 3 elements. 1 - edge_type, 2- host span of size_t, 3 - fanout vector
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename label_t,
-          bool store_transposed,
-          bool multi_gpu>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>,
-           std::optional<rmm::device_uvector<edge_type_t>>,
-           std::optional<rmm::device_uvector<int32_t>>,
-           std::optional<rmm::device_uvector<label_t>>,
-           std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  raft::device_span<vertex_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement                           = true,
-  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
-  bool dedupe_sources                             = false,
-  bool do_expensive_check                         = false);
-#endif
-
 /**
  * @brief Biased Neighborhood Sampling.
  *
@@ -461,131 +353,6 @@ neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
-
-
-
-
-
-
-
-
-
-
-
-
-
-#if 0
-/* FIXME:
- There are two options to support heterogeneous fanout
- 1) Create a new function 'heterogeneous_biased_neighbor_sample' which will take as input
-      only heterogeneous fanout type. Drawback: code redundancy
- 2) Update 'biased_neighbor_sample' to support both fanout types
-*/
-/**
- * @brief Biased Neighborhood Sampling.
- *
- * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
- *
- * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
- * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
- * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
- * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
- * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
- * (optional) will be described below and is dependent upon the input parameters.
- *
- * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
- * label and offsets values in the return set will be std::nullopt.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
- * the label output has values.  This will also result in the output being sorted by vertex label.
- * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
- * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
- * label output has values.  This will also result in the output being sorted by vertex label.  The
- * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
- * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
- * be shuffled so that all data with a particular label will be on the specified rank.
- *
- * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
- * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
- * @tparam weight_t Type of edge weights. Needs to be a floating point type.
- * @tparam edge_type_t Type of edge type. Needs to be an integral type.
- * @tparam label_t Type of label. Needs to be an integral type.
- * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
- * true) are major indices
- * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
- * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
- * handles to various CUDA libraries) to run graph algorithms.
- * @param graph_view Graph View object to generate NBR Sampling on.
- * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
- * @param edge_id_view Optional view object holding edge ids for @p graph_view.
- * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
- * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
- * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected.
- * @param starting_vertices Device span of starting vertex IDs for the sampling.
- * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
- * for the sampling.
- * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
- * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
- * output rank.  The label span must be sorted in ascending order.
- * @param fan_out Tuple of host spans mapping each edge type to fanout values. Element 0
- * of the tuple defines the size of the fanout per edge type while element 1 defines the branching out
- * (fan-out) degree per edge type for each level.
- * @param rng_state A pre-initialized raft::RngState object for generating random numbers
- * @param return_hops boolean flag specifying if the hop information should be returned
- * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
- * DEFAULT)
- * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
- * multiple times with the same label, it will only be passed once (for each label) as a source
- * for the next hop.  Default is false.
- * @param with_replacement boolean flag specifying if random sampling is done with replacement
- * (true); or, without replacement (false); default = true;
- * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
- * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
- * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
- * optional int32_t hop, optional label_t label, optional size_t offsets)
- */
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename bias_t,
-          typename label_t,
-          bool store_transposed,
-          bool multi_gpu>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>,
-           std::optional<rmm::device_uvector<edge_type_t>>,
-           std::optional<rmm::device_uvector<int32_t>>,
-           std::optional<rmm::device_uvector<label_t>>,
-           std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
-  raft::device_span<vertex_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement                           = true,
-  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
-  bool dedupe_sources                             = false,
-  bool do_expensive_check                         = false);
-#endif
-
 /*
  * @brief renumber sampled edge list and compress to the (D)CSR|(D)CSC format.
  *

From 3a6aeb2508e8b35ca2561e61c0e5a0885cfcdce5 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 04:31:22 -0700
Subject: [PATCH 008/117] add flag differentiating between biased and uniform
 sampling

---
 cpp/include/cugraph_c/sampling_algorithms.h       |  1 +
 cpp/src/c_api/neighbor_sampling.cpp               | 15 +++++++++++++--
 .../_cugraph_c/sampling_algorithms.pxd            |  1 +
 .../pylibcugraph/pylibcugraph/neighbor_sample.pyx |  6 +++++-
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index fcdb89dee2d..13e96bc3ab2 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -527,6 +527,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
+  bool_t is_biased,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 212f4736d42..4be914efc05 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -70,6 +70,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
+  bool is_biased_{false};
   cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
@@ -86,6 +87,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_resource_handle_t const* handle,
     cugraph_rng_state_t* rng_state,
     cugraph_graph_t* graph,
+    bool is_biased,
     cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
     cugraph_type_erased_device_array_view_t const* start_vertex_labels,
@@ -100,6 +102,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      is_biased_(is_biased),
       edge_biases_(
         reinterpret_cast<cugraph::c_api::cugraph_edge_property_view_t const*>(edge_biases)),
       start_vertices_(
@@ -222,8 +225,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
           (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
           // FIXME: add bias flag to differentiate between bias and uniform neighbor sample
-          // bias ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt ***
-          (edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view()),
+          is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+          //(edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view()),
           raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
           (start_vertex_labels_ != nullptr)
             ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
@@ -1028,10 +1031,13 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
     CUGRAPH_INVALID_INPUT,
     "fan_out should be of type int",
     *error);
+  
+  bool is_biased = false;
 
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
+                                    is_biased,
                                     nullptr,
                                     start_vertices,
                                     start_vertex_labels,
@@ -1050,6 +1056,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
+  bool_t is_biased,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
@@ -1106,6 +1113,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
+                                    is_biased,
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_labels,
@@ -1183,9 +1191,12 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
+  bool is_biased = true;
+
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
+                                    is_biased,
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_labels,
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index c47621db191..97cd7ba468c 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -68,6 +68,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_resource_handle_t* handle,
         cugraph_rng_state_t* rng_state,
         cugraph_graph_t* graph,
+        bool_t is_biased,
         const cugraph_edge_property_view_t* edge_biases,
         const cugraph_type_erased_device_array_view_t* start_vertices,
         const cugraph_type_erased_device_array_view_t* start_vertex_labels,
diff --git a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
index d3f8bbae9da..6b50a6294da 100644
--- a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
@@ -88,6 +88,7 @@ def neighbor_sample(ResourceHandle resource_handle,
                     _GPUGraph input_graph,
                     start_list,
                     h_fan_out,
+                    bool_t is_biased,
                     bool_t with_replacement,
                     bool_t do_expensive_check,
                     with_edge_properties=False,
@@ -125,11 +126,13 @@ def neighbor_sample(ResourceHandle resource_handle,
         Device array containing the list of starting vertices for sampling.
 
     edge_biases: FIXE: update this - Create edge_biases of type
-    'cugraph_edge_property_view_t'
+    'cugraph_edge_property_view_t' - edge biases not yet supported.
 
     h_fan_out: numpy array type or tuple of numpy array type
         Device array containing the brancing out (fan-out) degrees per
         starting vertex for each hop level.
+    
+    is_biased: bool # FIXME: Update docstrings
 
     with_replacement: bool
         If true, sampling procedure is done with replacement (the same vertex
@@ -406,6 +409,7 @@ def neighbor_sample(ResourceHandle resource_handle,
         c_resource_handle_ptr,
         rng_state_ptr,
         c_graph_ptr,
+        is_biased,
         <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
         start_ptr,
         batch_id_ptr,

From d2f6467114f105b797df4ccaaa3355bd1fc828da Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 04:40:11 -0700
Subject: [PATCH 009/117] update docstrings and rename variable

---
 cpp/include/cugraph_c/sampling_algorithms.h   | 11 ++++++++--
 cpp/src/c_api/neighbor_sampling.cpp           | 22 +++++++++----------
 .../dask/sampling/uniform_neighbor_sample.py  |  6 ++---
 .../sampling/uniform_neighbor_sample.py       |  6 ++---
 .../_cugraph_c/sampling_algorithms.pxd        |  2 +-
 5 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 13e96bc3ab2..e7b431ff1f0 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -336,10 +336,17 @@ typedef struct {
  * @brief       Create heterogeneous fanout
  *
  * Input data will be stored in the heterogenous_fanout.
+ * 
+ * The fanout is going to be a CSR structure, the edge_type_offsets will define which range
+ * of the fanout array is associated with each edge type, the fanout will be the values of
+ * fanout for that hop/type. So for edge type k, fanout[edge_type_offsets[k]] will identify
+ * the fanout for hop 0 for edge type k.  fanout[edge_type_offsets[k] +1] will identify the
+ * fanout for hop 1, etc. edge_type_offsets[k+1] will mark the beginning of the fanout
+ * array for type k+1 (and the end of the fanout array for type k.
  *
  * @param [in]  handle         Handle for accessing resources
  * @param [in]  graph          Pointer to graph
- * @param [in]  edge_type_size Type erased array of edge type size
+ * @param [in]  edge_type_offsets Type erased array of edge type size
  * @param [in]  fanout         Type erased array of fanout values
  * @param [out] heterogeneous_fanout Opaque pointer to fanout_t
  * @param [out] error          Pointer to an error object storing details of any error.  Will
@@ -349,7 +356,7 @@ typedef struct {
 cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
-  const cugraph_type_erased_host_array_view_t* edge_type_size,
+  const cugraph_type_erased_host_array_view_t* edge_type_offsets,
   const cugraph_type_erased_host_array_view_t* fanout,
   cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
   cugraph_error_t** error);
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 4be914efc05..4f86fb4dbec 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -440,7 +440,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
   cugraph::c_api::cugraph_graph_t* graph_;
-  cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_size_;
+  cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_offsets_;
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
   // FIXME: This type doesn't exist: instead create an
   // 'std::tuple<cugraph_type_erased_host_array_t*>'
@@ -448,14 +448,14 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
 
   create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
                                       ::cugraph_graph_t* graph,
-                                      ::cugraph_type_erased_host_array_view_t const* edge_type_size,
+                                      ::cugraph_type_erased_host_array_view_t const* edge_type_offsets,
                                       ::cugraph_type_erased_host_array_view_t const* fanout)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-      edge_type_size_(
+      edge_type_offsets_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-          edge_type_size)),
+          edge_type_offsets)),
       fanout_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
   {
@@ -473,12 +473,12 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
     if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
       unsupported();
     } else {
-      std::vector<int32_t> edge_type_size_copy{(int32_t)edge_type_size_->size_};
+      std::vector<int32_t> edge_type_offsets_copy{(int32_t)edge_type_offsets_->size_};
       std::vector<int32_t> fanout_copy{(int32_t)fanout_->size_};
 
-      raft::copy(edge_type_size_copy.data(),
-                 edge_type_size_->as_type<int32_t>(),
-                 edge_type_size_->size_,
+      raft::copy(edge_type_offsets_copy.data(),
+                 edge_type_offsets_->as_type<int32_t>(),
+                 edge_type_offsets_->size_,
                  handle_.get_stream());
 
       raft::copy(
@@ -487,7 +487,7 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
       // std::tuple (template)
       // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
       result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
-        new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_size_copy, INT32),
+        new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_offsets_copy, INT32),
         new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)};
     }
   }
@@ -1213,12 +1213,12 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
 extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
-  const cugraph_type_erased_host_array_view_t* edge_type_size,
+  const cugraph_type_erased_host_array_view_t* edge_type_offsets,
   const cugraph_type_erased_host_array_view_t* fanout,
   cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
   cugraph_error_t** error)
 {
-  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_size, fanout);
+  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_offsets, fanout);
 
   return cugraph::c_api::run_algorithm(graph, functor, heterogeneous_fanout, error);
 }
diff --git a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index af736ce7846..89954a34ad4 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -617,11 +617,11 @@ def uniform_neighbor_sample(
     elif isinstance(fanout_vals, dict):
         # FIXME: Add expensive check to ensure all dict values are lists
         # Convert to a tuple of sequence (edge type size and fanout values)
-        edge_type_size = []
-        [edge_type_size.append(len(s)) for s in list(fanout_vals.values())]
+        edge_type_offsets = []
+        [edge_type_offsets.append(len(s)) for s in list(fanout_vals.values())]
         edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
         fanout_vals = (
-            numpy.asarray(edge_type_size, dtype="int32"),
+            numpy.asarray(edge_type_offsets, dtype="int32"),
             numpy.asarray(edge_type_fanout_vals, dtype="int32"),
         )
     else:
diff --git a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
index d17e070ad3a..5083d50e7fb 100644
--- a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
@@ -319,11 +319,11 @@ def uniform_neighbor_sample(
     elif isinstance(fanout_vals, dict):
         # FIXME: Add expensive check to ensure all dict values are lists
         # Convert to a tuple of sequence (edge type size and fanout values)
-        edge_type_size = []
-        [edge_type_size.append(len(s)) for s in list(fanout_vals.values())]
+        edge_type_offsets = []
+        [edge_type_offsets.append(len(s)) for s in list(fanout_vals.values())]
         edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
         fanout_vals = (
-            numpy.asarray(edge_type_size, dtype="int32"),
+            numpy.asarray(edge_type_offsets, dtype="int32"),
             numpy.asarray(edge_type_fanout_vals, dtype="int32"),
         )
     else:
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 97cd7ba468c..9bcbaa67dfb 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -131,7 +131,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_create_heterogeneous_fanout(
             const cugraph_resource_handle_t* handle,
             cugraph_graph_t* graph,
-            const cugraph_type_erased_host_array_view_t* edge_type_size,
+            const cugraph_type_erased_host_array_view_t* edge_type_offsets,
             const cugraph_type_erased_host_array_view_t* fanout,
             cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
             cugraph_error_t** error

From 5d25155e14921b62401f77ceddab57b41e37cae4 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 04:51:30 -0700
Subject: [PATCH 010/117] rename variable

---
 cpp/include/cugraph_c/sampling_algorithms.h | 2 +-
 cpp/src/c_api/neighbor_sampling.cpp         | 4 ++--
 cpp/src/c_api/neighbor_sampling.hpp         | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index e7b431ff1f0..d55a6fd9b55 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -346,7 +346,7 @@ typedef struct {
  *
  * @param [in]  handle         Handle for accessing resources
  * @param [in]  graph          Pointer to graph
- * @param [in]  edge_type_offsets Type erased array of edge type size
+ * @param [in]  edge_type_offsets Type erased array of edge type offsets
  * @param [in]  fanout         Type erased array of fanout values
  * @param [out] heterogeneous_fanout Opaque pointer to fanout_t
  * @param [out] error          Pointer to an error object storing details of any error.  Will
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 4f86fb4dbec..6ed199e9725 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -245,8 +245,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
           (heterogeneous_fan_out_ != nullptr)
             ? std::make_optional(std::make_tuple(
-                raft::host_span<const int>{heterogeneous_fan_out_->edge_type_id->as_type<int>(),
-                                           heterogeneous_fan_out_->edge_type_id->size_},
+                raft::host_span<const int>{heterogeneous_fan_out_->edge_type_offsets->as_type<int>(),
+                                           heterogeneous_fan_out_->edge_type_offsets->size_},
                 raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
                                            heterogeneous_fan_out_->fanout->size_}))
             : std::nullopt,
diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
index 346ad150fab..7c4ab839813 100644
--- a/cpp/src/c_api/neighbor_sampling.hpp
+++ b/cpp/src/c_api/neighbor_sampling.hpp
@@ -22,7 +22,7 @@ namespace cugraph {
 namespace c_api {
 
 struct cugraph_sample_heterogeneous_fanout_t {
-  cugraph_type_erased_host_array_t* edge_type_id;
+  cugraph_type_erased_host_array_t* edge_type_offsets;
   cugraph_type_erased_host_array_t* fanout;
 };
 

From 80f8b8692acd16afc75f60369f5982de56b7ac81 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 07:48:43 -0700
Subject: [PATCH 011/117] create new tuple type

---
 cpp/include/cugraph/sampling_functions.hpp    |  7 +-
 cpp/include/cugraph_c/sampling_algorithms.h   |  5 +-
 cpp/src/c_api/graph_functions.cpp             |  2 +-
 cpp/src/c_api/neighbor_sampling.cpp           | 38 ++++++----
 cpp/src/c_api/neighbor_sampling.hpp           |  6 +-
 .../_cugraph_c/sampling_algorithms.pxd        |  1 -
 .../pylibcugraph/uniform_neighbor_sample.pyx  | 75 ++++---------------
 7 files changed, 52 insertions(+), 82 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index aee0c7f240e..d0a91ca04e1 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -42,7 +42,8 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
 
 /**
  * @brief Uniform Neighborhood Sampling.
- * # FIXME: Reset all changes to uniform neighborhood sampling
+ * 
+ * @deprecated  This API will be deleted, use neighbor_sample instead
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects from these outgoing neighbors to extract a subgraph.
@@ -131,8 +132,6 @@ uniform_neighbor_sample(
   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::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement                           = true,
@@ -142,6 +141,8 @@ uniform_neighbor_sample(
 
 /**
  * @brief Biased Neighborhood Sampling.
+ * 
+ * @deprecated  This API will be deleted, use neighbor_sample instead
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index d55a6fd9b55..b5a16668c76 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -370,6 +370,8 @@ void cugraph_heterogeneous_fanout_free(cugraph_sample_heterogeneous_fanout_t* he
 
 /**
  * @brief     Uniform Neighborhood Sampling
+ * 
+ * @deprecated  This API will be deleted, use cugraph_neighbor_sample instead
  *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
@@ -417,7 +419,6 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -426,6 +427,8 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
 
 /**
  * @brief     Biased Neighborhood Sampling
+ * 
+ * @deprecated  This API will be deleted, use cugraph_neighbor_sample instead
  *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
diff --git a/cpp/src/c_api/graph_functions.cpp b/cpp/src/c_api/graph_functions.cpp
index 349598b7ef6..8778369dbe6 100644
--- a/cpp/src/c_api/graph_functions.cpp
+++ b/cpp/src/c_api/graph_functions.cpp
@@ -84,7 +84,7 @@ struct create_vertex_pairs_functor : public cugraph::c_api::abstract_functor {
                               std::nullopt,
                               std::nullopt);
       }
-      // std::tuple (template)
+      // FIXME: use std::tuple (template) instead.
       result_ = new cugraph::c_api::cugraph_vertex_pairs_t{
         new cugraph::c_api::cugraph_type_erased_device_array_t(first_copy, graph_->vertex_type_),
         new cugraph::c_api::cugraph_type_erased_device_array_t(second_copy, graph_->vertex_type_)};
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 6ed199e9725..e1866caece0 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -121,9 +121,11 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           label_offsets)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
+
       heterogeneous_fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(
           heterogeneous_fan_out)),
+  
       options_(options),
       do_expensive_check_(do_expensive_check)
   {
@@ -215,6 +217,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_first(),
         graph_view.local_vertex_partition_range_last(),
         do_expensive_check_);
+      
+      // FIXME: Consolidate 'fan_out_' and 'heterogeneous_fan_out_' into one
+      // argument with std::variant
 
       auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
         cugraph::neighbor_sample(
@@ -224,9 +229,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
           (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
           (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-          // FIXME: add bias flag to differentiate between bias and uniform neighbor sample
           is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-          //(edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view()),
           raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
           (start_vertex_labels_ != nullptr)
             ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
@@ -242,7 +245,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           (fan_out_ != nullptr) ? std::make_optional<raft::host_span<const int>>(
                                     fan_out_->as_type<const int>(), fan_out_->size_)
                                 : std::nullopt,
-
+          /*
           (heterogeneous_fan_out_ != nullptr)
             ? std::make_optional(std::make_tuple(
                 raft::host_span<const int>{heterogeneous_fan_out_->edge_type_offsets->as_type<int>(),
@@ -250,6 +253,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                 raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
                                            heterogeneous_fan_out_->fanout->size_}))
             : std::nullopt,
+          */
+          std::nullopt,
           options_.return_hops_,
           options_.with_replacement_,
           options_.prior_sources_behavior_,
@@ -442,10 +447,7 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
   cugraph::c_api::cugraph_graph_t* graph_;
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_offsets_;
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
-  // FIXME: This type doesn't exist: instead create an
-  // 'std::tuple<cugraph_type_erased_host_array_t*>'
   cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
-
   create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
                                       ::cugraph_graph_t* graph,
                                       ::cugraph_type_erased_host_array_view_t const* edge_type_offsets,
@@ -484,11 +486,14 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
       raft::copy(
         fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
 
-      // std::tuple (template)
-      // result_ = new std::tuple <template type of 2 cugraph_type_erased_host_array_t>
-      result_ = new cugraph::c_api::cugraph_sample_heterogeneous_fanout_t{
+      auto result_tuple = std::make_tuple(
         new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_offsets_copy, INT32),
-        new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)};
+        new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)
+      );
+
+      result_ = &result_tuple;
+  
+
     }
   }
 };
@@ -985,7 +990,6 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -1045,7 +1049,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                     label_to_comm_rank,
                                     label_offsets,
                                     fan_out,
-                                    heterogeneous_fanout,
+                                    nullptr,
                                     std::move(options_cpp),
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
@@ -1072,6 +1076,15 @@ cugraph_error_code_t cugraph_neighbor_sample(
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
+  if (is_biased) {
+    CAPI_EXPECTS(
+      (edge_biases != nullptr) ||
+        (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
+      CUGRAPH_INVALID_INPUT,
+      "edge_biases is required if the graph is not weighted",
+      *error);
+  }
+  
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify label_offsets if retain_seeds is true",
@@ -1145,7 +1158,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
-  // FIXME: Check if bias = True before checking the statement below
   CAPI_EXPECTS(
     (edge_biases != nullptr) ||
       (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
index 7c4ab839813..9c82f191f3a 100644
--- a/cpp/src/c_api/neighbor_sampling.hpp
+++ b/cpp/src/c_api/neighbor_sampling.hpp
@@ -21,10 +21,8 @@
 namespace cugraph {
 namespace c_api {
 
-struct cugraph_sample_heterogeneous_fanout_t {
-  cugraph_type_erased_host_array_t* edge_type_offsets;
-  cugraph_type_erased_host_array_t* fanout;
-};
+using cugraph_sample_heterogeneous_fanout_t = std::tuple<
+  cugraph_type_erased_host_array_t*, cugraph_type_erased_host_array_t*>;
 
 }  // namespace c_api
 }  // namespace cugraph
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 9bcbaa67dfb..4b0252b614c 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -56,7 +56,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
         const cugraph_type_erased_host_array_view_t* fan_out,
-        const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
         cugraph_rng_state_t* rng_state,
         const cugraph_sampling_options_t* options,
         bool_t do_expensive_check,
diff --git a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
index c4887414122..f002622f497 100644
--- a/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/uniform_neighbor_sample.pyx
@@ -53,9 +53,7 @@ from pylibcugraph._cugraph_c.algorithms cimport (
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
     cugraph_uniform_neighbor_sample,
-    cugraph_sample_heterogeneous_fanout_t,
-    cugraph_create_heterogeneous_fanout,
-    cugraph_heterogeneous_fanout_free,
+
 )
 from pylibcugraph.resource_handle cimport (
     ResourceHandle,
@@ -85,6 +83,7 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
                             _GPUGraph input_graph,
                             start_list,
                             h_fan_out,
+                            *,
                             bool_t with_replacement,
                             bool_t do_expensive_check,
                             with_edge_properties=False,
@@ -117,7 +116,7 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     start_list: device array type
         Device array containing the list of starting vertices for sampling.
 
-    h_fan_out: numpy array type or tuple of numpy array type
+    h_fan_out: numpy array type
         Device array containing the brancing out (fan-out) degrees per
         starting vertex for each hop level.
 
@@ -212,70 +211,22 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     )
 
     cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
-    cdef cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fanout_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
 
     cdef bool_t c_deduplicate_sources = deduplicate_sources
     cdef bool_t c_return_hops = return_hops
     cdef bool_t c_renumber = renumber
     cdef bool_t c_compress_per_hop = compress_per_hop
 
-    cdef cugraph_error_code_t error_code
-    cdef cugraph_error_t* error_ptr
-    cdef uintptr_t ai_fan_out_ptr
-    cdef uintptr_t ai_fan_out_size_ptr
-    cdef uintptr_t ai_fan_out_values_ptr
-
-
     assert_CAI_type(start_list, "start_list")
     assert_CAI_type(batch_id_list, "batch_id_list", True)
     assert_CAI_type(label_list, "label_list", True)
     assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
     assert_CAI_type(label_offsets, "label_offsets", True)
-    if not isinstance(h_fan_out, tuple):
-        assert_AI_type(h_fan_out, "h_fan_out")
-        ai_fan_out_ptr = \
-            h_fan_out.__array_interface__["data"][0]
-
-        fan_out_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_ptr,
-                len(h_fan_out),
-                get_c_type_from_numpy_type(h_fan_out.dtype))
-    else:
-        assert_AI_type(h_fan_out[0], "h_fan_out_size")
-        assert_AI_type(h_fan_out[1], "h_fan_out_values")
-        ai_fan_out_size_ptr = \
-            h_fan_out[0].__array_interface__["data"][0]
-        ai_fan_out_values_ptr = \
-            h_fan_out[1].__array_interface__["data"][0]
-
-        fan_out_size_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_size_ptr,
-                len(h_fan_out[0]),
-                get_c_type_from_numpy_type(h_fan_out[0].dtype))
-
-        fan_out_values_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_values_ptr,
-                len(h_fan_out[1]),
-                get_c_type_from_numpy_type(h_fan_out[1].dtype))
-
-        error_code = cugraph_create_heterogeneous_fanout(
-            c_resource_handle_ptr,
-            c_graph_ptr,
-            fan_out_size_ptr,
-            fan_out_values_ptr,
-            &heterogeneous_fanout_ptr,
-            &error_ptr
-        )
-
-    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fanout")
+    assert_AI_type(h_fan_out, "h_fan_out")
 
     cdef cugraph_sample_result_t* result_ptr
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
 
     cdef uintptr_t cai_start_ptr = \
         start_list.__cuda_array_interface__["data"][0]
@@ -300,6 +251,9 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
         cai_label_offsets_ptr = \
             label_offsets.__cuda_array_interface__['data'][0]
 
+    cdef uintptr_t ai_fan_out_ptr = \
+        h_fan_out.__array_interface__["data"][0]
+
     cdef cugraph_type_erased_device_array_view_t* start_ptr = \
         cugraph_type_erased_device_array_view_create(
             <void*>cai_start_ptr,
@@ -344,6 +298,13 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
                 get_c_type_from_numpy_type(label_offsets.dtype)
             )
 
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_ptr,
+            len(h_fan_out),
+            get_c_type_from_numpy_type(h_fan_out.dtype))
+
+
     cg_rng_state = CuGraphRandomState(resource_handle, random_state)
 
     cdef cugraph_rng_state_t* rng_state_ptr = \
@@ -401,7 +362,6 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
         label_to_output_comm_rank_ptr,
         label_offsets_ptr,
         fan_out_ptr,
-        heterogeneous_fanout_ptr,
         rng_state_ptr,
         sampling_options,
         do_expensive_check,
@@ -412,9 +372,6 @@ def uniform_neighbor_sample(ResourceHandle resource_handle,
     # Free the sampling options
     cugraph_sampling_options_free(sampling_options)
 
-    if isinstance(h_fan_out, tuple):
-        cugraph_heterogeneous_fanout_free(heterogeneous_fanout_ptr)
-
     # Free the two input arrays that are no longer needed.
     cugraph_type_erased_device_array_view_free(start_ptr)
     cugraph_type_erased_host_array_view_free(fan_out_ptr)

From 50e0fc561cfa2c3f85bcba527c01628553e49f58 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 07:52:15 -0700
Subject: [PATCH 012/117] remove unnecessary check

---
 cpp/src/c_api/neighbor_sampling.cpp | 32 ++++++++++++-----------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index e1866caece0..de6adfc9cb0 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -471,30 +471,24 @@ struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_fun
             bool multi_gpu>
   void operator()()
   {
-    // FIXME: Remove this check as it is not necessary
-    if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
-      unsupported();
-    } else {
-      std::vector<int32_t> edge_type_offsets_copy{(int32_t)edge_type_offsets_->size_};
-      std::vector<int32_t> fanout_copy{(int32_t)fanout_->size_};
 
-      raft::copy(edge_type_offsets_copy.data(),
-                 edge_type_offsets_->as_type<int32_t>(),
-                 edge_type_offsets_->size_,
-                 handle_.get_stream());
+    std::vector<int32_t> edge_type_offsets_copy{(int32_t)edge_type_offsets_->size_};
+    std::vector<int32_t> fanout_copy{(int32_t)fanout_->size_};
 
-      raft::copy(
-        fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
+    raft::copy(edge_type_offsets_copy.data(),
+                edge_type_offsets_->as_type<int32_t>(),
+                edge_type_offsets_->size_,
+                handle_.get_stream());
 
-      auto result_tuple = std::make_tuple(
-        new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_offsets_copy, INT32),
-        new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)
-      );
+    raft::copy(
+      fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
 
-      result_ = &result_tuple;
-  
+    auto result_tuple = std::make_tuple(
+      new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_offsets_copy, INT32),
+      new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)
+    );
 
-    }
+    result_ = &result_tuple;
   }
 };
 }  // namespace

From 9f455bfc4fa3dd5dc88e3aec951333a5c2847def Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 10:09:07 -0700
Subject: [PATCH 013/117] add constructor converting from array_view_t to
 array_t

---
 cpp/src/c_api/array.hpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/cpp/src/c_api/array.hpp b/cpp/src/c_api/array.hpp
index 2dce70fa924..2fa58d917ae 100644
--- a/cpp/src/c_api/array.hpp
+++ b/cpp/src/c_api/array.hpp
@@ -125,6 +125,15 @@ struct cugraph_type_erased_host_array_t {
     std::copy(vec.begin(), vec.end(), reinterpret_cast<T*>(data_.get()));
   }
 
+  cugraph_type_erased_host_array_t(cugraph_type_erased_host_array_view_t const* view_p)
+    : data_(std::make_unique<std::byte[]>(view_p->num_bytes_)),
+      size_(view_p->size_),
+      num_bytes_(view_p->num_bytes_),
+      type_(view_p->type_)
+  {
+    std::copy(view_p->data_, view_p->data_ + num_bytes_, data_.get());
+  }
+
   template <typename T>
   T* as_type()
   {

From d114534065056bc8c8e74d44455905fbb13cd5fc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 10:09:35 -0700
Subject: [PATCH 014/117] leverage new constructor and remove unnecessary code

---
 cpp/src/c_api/neighbor_sampling.cpp | 73 ++++++-----------------------
 1 file changed, 14 insertions(+), 59 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index de6adfc9cb0..d37801b7b76 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -221,6 +221,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       // FIXME: Consolidate 'fan_out_' and 'heterogeneous_fan_out_' into one
       // argument with std::variant
 
+      //cugraph_type_erased_host_array_t* x = std::get<0>(heterogeneous_fan_out_);
+
+      //cugraph::c_api::cugraph_type_erased_host_array_t* x = std::get<0>(*heterogeneous_fan_out_);
+
       auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
         cugraph::neighbor_sample(
           handle_,
@@ -245,16 +249,14 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           (fan_out_ != nullptr) ? std::make_optional<raft::host_span<const int>>(
                                     fan_out_->as_type<const int>(), fan_out_->size_)
                                 : std::nullopt,
-          /*
+          
           (heterogeneous_fan_out_ != nullptr)
             ? std::make_optional(std::make_tuple(
-                raft::host_span<const int>{heterogeneous_fan_out_->edge_type_offsets->as_type<int>(),
-                                           heterogeneous_fan_out_->edge_type_offsets->size_},
-                raft::host_span<const int>{heterogeneous_fan_out_->fanout->as_type<int>(),
-                                           heterogeneous_fan_out_->fanout->size_}))
+                raft::host_span<const int>{std::get<0>(*heterogeneous_fan_out_)->as_type<int>(),
+                                           std::get<0>(*heterogeneous_fan_out_)->size_},
+                raft::host_span<const int>{std::get<1>(*heterogeneous_fan_out_)->as_type<int>(),
+                                           std::get<1>(*heterogeneous_fan_out_)->size_}))
             : std::nullopt,
-          */
-          std::nullopt,
           options_.return_hops_,
           options_.with_replacement_,
           options_.prior_sources_behavior_,
@@ -441,56 +443,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   }
 };
 
-
-struct create_heterogeneous_fanout_functor : public cugraph::c_api::abstract_functor {
-  raft::handle_t const& handle_;
-  cugraph::c_api::cugraph_graph_t* graph_;
-  cugraph::c_api::cugraph_type_erased_host_array_view_t const* edge_type_offsets_;
-  cugraph::c_api::cugraph_type_erased_host_array_view_t const* fanout_;
-  cugraph::c_api::cugraph_sample_heterogeneous_fanout_t* result_{};
-  create_heterogeneous_fanout_functor(::cugraph_resource_handle_t const* handle,
-                                      ::cugraph_graph_t* graph,
-                                      ::cugraph_type_erased_host_array_view_t const* edge_type_offsets,
-                                      ::cugraph_type_erased_host_array_view_t const* fanout)
-    : abstract_functor(),
-      handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
-      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-      edge_type_offsets_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-          edge_type_offsets)),
-      fanout_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))
-  {
-  }
-
-  template <typename vertex_t,
-            typename edge_t,
-            typename weight_t,
-            typename edge_type_type_t,
-            bool store_transposed,
-            bool multi_gpu>
-  void operator()()
-  {
-
-    std::vector<int32_t> edge_type_offsets_copy{(int32_t)edge_type_offsets_->size_};
-    std::vector<int32_t> fanout_copy{(int32_t)fanout_->size_};
-
-    raft::copy(edge_type_offsets_copy.data(),
-                edge_type_offsets_->as_type<int32_t>(),
-                edge_type_offsets_->size_,
-                handle_.get_stream());
-
-    raft::copy(
-      fanout_copy.data(), fanout_->as_type<int32_t>(), fanout_->size_, handle_.get_stream());
-
-    auto result_tuple = std::make_tuple(
-      new cugraph::c_api::cugraph_type_erased_host_array_t(edge_type_offsets_copy, INT32),
-      new cugraph::c_api::cugraph_type_erased_host_array_t(fanout_copy, INT32)
-    );
-
-    result_ = &result_tuple;
-  }
-};
 }  // namespace
 
 extern "C" cugraph_error_code_t cugraph_sampling_options_create(
@@ -1224,7 +1176,10 @@ extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
   cugraph_error_t** error)
 {
-  create_heterogeneous_fanout_functor functor(handle, graph, edge_type_offsets, fanout);
 
-  return cugraph::c_api::run_algorithm(graph, functor, heterogeneous_fanout, error);
+  *heterogeneous_fanout = reinterpret_cast<cugraph_sample_heterogeneous_fanout_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
+          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_offsets)),
+          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))});
+  
+  return CUGRAPH_SUCCESS;
 }

From cf4a3ae36187f78429e67b9304df70e2fb3602c4 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 10:57:00 -0700
Subject: [PATCH 015/117] ensure edge types are ordered in increasing order

---
 .../cugraph/dask/sampling/uniform_neighbor_sample.py     | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index 89954a34ad4..45805b88581 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 from itertools import chain
+import collections
+from itertools import accumulate
 
 import warnings
 
@@ -617,9 +619,16 @@ def uniform_neighbor_sample(
     elif isinstance(fanout_vals, dict):
         # FIXME: Add expensive check to ensure all dict values are lists
         # Convert to a tuple of sequence (edge type size and fanout values)
+
+        fanout_vals = {2:[1], 3:[4, 5], 1:[6, 7], 0:[8, 9, 10]}
+        fanout_vals = collections.OrderedDict(sorted(fanout_vals.items()))
+
         edge_type_offsets = []
         [edge_type_offsets.append(len(s)) for s in list(fanout_vals.values())]
+
+        edge_type_offsets = list(accumulate(edge_type_offsets))
         edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
+
         fanout_vals = (
             numpy.asarray(edge_type_offsets, dtype="int32"),
             numpy.asarray(edge_type_fanout_vals, dtype="int32"),

From bc87b501b5461fef76f54f4475d26bbf02b069e3 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 10:57:27 -0700
Subject: [PATCH 016/117] update docstrings

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index b82e306cc4b..9adab8ad7ed 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -85,17 +85,12 @@ neighbor_sample_impl(
       "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                                 // to store hops
   } else {
-    /*
-    FIXME: Since we are dealing with offsets and size, just check that the last element of the 
-    first host array == size of the second array.
 
     CUGRAPH_EXPECTS(
-      std::accumulate(std::get<0>(*heterogeneous_fan_out).begin(),
-                      std::get<0>(*heterogeneous_fan_out).end(),
-                      0) == std::get<1>(*heterogeneous_fan_out).size() &&
+      std::get<0>(*heterogeneous_fan_out).back() == std::get<1>(*heterogeneous_fan_out).size() &&
         std::get<1>(*heterogeneous_fan_out).size() != 0,
       "Invalid input argument: number of levels and size must match and should be non zero.");
-    */
+    
 
     CUGRAPH_EXPECTS(
       std::get<0>(*heterogeneous_fan_out).size() <=

From 301368483c49ce94934deea90a91da2871dc4f98 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 21 Aug 2024 12:20:49 -0700
Subject: [PATCH 017/117] update docstrings

---
 cpp/src/c_api/neighbor_sampling.cpp | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index d37801b7b76..6190b1a3c73 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -220,11 +220,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       
       // FIXME: Consolidate 'fan_out_' and 'heterogeneous_fan_out_' into one
       // argument with std::variant
-
-      //cugraph_type_erased_host_array_t* x = std::get<0>(heterogeneous_fan_out_);
-
-      //cugraph::c_api::cugraph_type_erased_host_array_t* x = std::get<0>(*heterogeneous_fan_out_);
-
       auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
         cugraph::neighbor_sample(
           handle_,

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 018/117] 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,

From 068b0a387d2d5988d20423edc8cc75a7f5d87d84 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 12:42:35 -0700
Subject: [PATCH 019/117] undo changes to uniform neighbor sample

---
 cpp/include/cugraph/sampling_functions.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index d0a91ca04e1..2cf159b1c78 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -131,7 +131,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,
+  raft::host_span<int32_t const> fan_out,
   raft::random::RngState& rng_state,
   bool return_hops,
   bool with_replacement                           = true,

From 6920f65ac2422fed160a8f9585ba51d73ca8c47f Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 14:52:33 -0700
Subject: [PATCH 020/117] update docstrings

---
 cpp/include/cugraph/sampling_functions.hpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 2cf159b1c78..3140e9073a0 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -250,6 +250,8 @@ biased_neighbor_sample(
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
+ * When branching out to select outgoing neighbors, either fan_out or heterogeneous_fan_out must
+ * be provided but not both.
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
  * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
@@ -290,7 +292,8 @@ biased_neighbor_sample(
  * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
  * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
  * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected.
+ * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
+ * uniform sampling.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
  * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
@@ -299,11 +302,13 @@ biased_neighbor_sample(
  * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
  * output rank.  The label span must be sorted in ascending order.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level
+ * level. When fan_out is provided, the sampling method uses the same fanout value for each type.
  * @param heterogeneous_fan_out Tuple of host spans defining branching out (fan-out) degree per
  * source vertex for each level in CSR style format. The first element of the tuple is the offset
  * array per edge type id and the second element correspond to the fanout values.
- * @param return_hops boolean flag specifying if the hop information should be returned
+ * When heterogeneous_fan_out is provided, different fan_out values can be used for each edge type.
+ * The fan-out offsets size must be proportional to the number of edge types and fan_out values.
+ * @param return_hops boolean flag specifying if the hop information should be returned. 
  * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
  * DEFAULT)
  * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X

From 760c5cd28fb1d44a4109fc7b6c8ffca6247d936a Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 15:11:52 -0700
Subject: [PATCH 021/117] re-order arguments

---
 cpp/include/cugraph_c/sampling_algorithms.h        |  5 ++++-
 cpp/src/c_api/neighbor_sampling.cpp                | 14 +++++++-------
 .../_cugraph_c/sampling_algorithms.pxd             |  2 +-
 .../pylibcugraph/pylibcugraph/neighbor_sample.pyx  |  2 +-
 4 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index b5a16668c76..d117a329510 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -526,6 +526,9 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * We only support type INT32 for both the offsets and the fanout values array.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
+ * @param [in]  is_biased
+ *                           A flag specifying whether to run biased neighborhood sampling
+ *                           (if set to true) or uniform neighbor sampling.
  * @param [in]  do_expensive_check
  *                           A flag to run expensive checks for input arguments (if set to true)
  * @param [out]  result      Output from the uniform_neighbor_sample call
@@ -537,7 +540,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
-  bool_t is_biased,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
@@ -547,6 +549,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   const cugraph_sampling_options_t* options,
+  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 6190b1a3c73..fe0212bb4cd 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -70,7 +70,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
-  bool is_biased_{false};
   cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
@@ -80,6 +79,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
+  bool is_biased_{false};
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
@@ -87,7 +87,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_resource_handle_t const* handle,
     cugraph_rng_state_t* rng_state,
     cugraph_graph_t* graph,
-    bool is_biased,
     cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
     cugraph_type_erased_device_array_view_t const* start_vertex_labels,
@@ -97,12 +96,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_type_erased_host_array_view_t const* fan_out,
     cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out,
     cugraph::c_api::cugraph_sampling_options_t options,
+    bool is_biased,
     bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
       graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
-      is_biased_(is_biased),
       edge_biases_(
         reinterpret_cast<cugraph::c_api::cugraph_edge_property_view_t const*>(edge_biases)),
       start_vertices_(
@@ -127,6 +126,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           heterogeneous_fan_out)),
   
       options_(options),
+      is_biased_(is_biased),
       do_expensive_check_(do_expensive_check)
   {
   }
@@ -982,7 +982,6 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
-                                    is_biased,
                                     nullptr,
                                     start_vertices,
                                     start_vertex_labels,
@@ -992,6 +991,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                     fan_out,
                                     nullptr,
                                     std::move(options_cpp),
+                                    is_biased,
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1001,7 +1001,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
-  bool_t is_biased,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_labels,
@@ -1011,6 +1010,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
   const cugraph_sampling_options_t* options,
+  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
@@ -1067,7 +1067,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
-                                    is_biased,
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_labels,
@@ -1077,6 +1076,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
                                     fan_out,
                                     heterogeneous_fanout,
                                     std::move(options_cpp),
+                                    is_biased,
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1149,7 +1149,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
-                                    is_biased,
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_labels,
@@ -1159,6 +1158,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                     fan_out,
                                     nullptr,
                                     std::move(options_cpp),
+                                    is_biased,
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 4b0252b614c..2074e88725f 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -67,7 +67,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_resource_handle_t* handle,
         cugraph_rng_state_t* rng_state,
         cugraph_graph_t* graph,
-        bool_t is_biased,
         const cugraph_edge_property_view_t* edge_biases,
         const cugraph_type_erased_device_array_view_t* start_vertices,
         const cugraph_type_erased_device_array_view_t* start_vertex_labels,
@@ -77,6 +76,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
         const cugraph_sampling_options_t* options,
+        bool_t is_biased,
         bool_t do_expensive_check,
         cugraph_sample_result_t** result,
         cugraph_error_t** error
diff --git a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
index 6b50a6294da..e8ce3e26992 100644
--- a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
@@ -409,7 +409,6 @@ def neighbor_sample(ResourceHandle resource_handle,
         c_resource_handle_ptr,
         rng_state_ptr,
         c_graph_ptr,
-        is_biased,
         <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
         start_ptr,
         batch_id_ptr,
@@ -419,6 +418,7 @@ def neighbor_sample(ResourceHandle resource_handle,
         fan_out_ptr,
         heterogeneous_fanout_ptr,
         sampling_options,
+        is_biased,
         do_expensive_check,
         &result_ptr,
         &error_ptr)

From 1e0ef27ee90fc0a030bdf13dcb60c82b0aad098d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 15:25:16 -0700
Subject: [PATCH 022/117] remove outdated comments

---
 cpp/src/c_api/neighbor_sampling.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index fe0212bb4cd..90dfd6084f9 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "c_api/neighbor_sampling.hpp"  // FIXME: Remove this and instead use std::tuple
+#include "c_api/neighbor_sampling.hpp"
 
 #include "c_api/abstract_functor.hpp"
 #include "c_api/graph.hpp"
@@ -996,7 +996,6 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-// FIXME: Add flag for bias=True/False
 cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,

From de7962098e9b098bcb9a35bf19045a8504f62775 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 17:43:02 -0700
Subject: [PATCH 023/117] add arguments and type check

---
 cpp/src/c_api/neighbor_sampling.cpp | 48 +++++++++++++++++++++--------
 1 file changed, 36 insertions(+), 12 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 90dfd6084f9..019692829c4 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -1007,7 +1007,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fan_out,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
@@ -1047,6 +1047,36 @@ cugraph_error_code_t cugraph_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "cannot specify label_to_comm_rank unless label_list is also specified",
                *error);
+  
+  CAPI_EXPECTS(!((fan_out != nullptr) && (heterogeneous_fan_out != nullptr)),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify both fan_out and heterogeneous_fan_out",
+               *error);
+  
+  if (fan_out != nullptr) {
+     CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                   fan_out)
+                   ->type_ == INT32,
+               CUGRAPH_INVALID_INPUT,
+               "fan_out type must be INT32",
+               *error);
+    
+  } else {
+    
+    CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)))
+                   ->type_ == INT32,
+               CUGRAPH_INVALID_INPUT,
+               "edge type offsets type must be INT32",
+               *error);
+    
+    CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)))
+                   ->type_ == INT32,
+               CUGRAPH_INVALID_INPUT,
+               "fan_out values type must be INT32",
+               *error);
+  }
 
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -1056,12 +1086,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
                "vertex type of graph and start_vertices must match",
                *error);
 
-  CAPI_EXPECTS(
-    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
-        ->type_ == INT32,
-    CUGRAPH_INVALID_INPUT,
-    "fan_out should be of type int",
-    *error);
 
   neighbor_sampling_functor functor{handle,
                                     rng_state,
@@ -1073,7 +1097,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
                                     label_to_comm_rank,
                                     label_offsets,
                                     fan_out,
-                                    heterogeneous_fanout,
+                                    heterogeneous_fan_out,
                                     std::move(options_cpp),
                                     is_biased,
                                     do_expensive_check};
@@ -1166,14 +1190,14 @@ extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
   const cugraph_type_erased_host_array_view_t* edge_type_offsets,
-  const cugraph_type_erased_host_array_view_t* fanout,
-  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fan_out,
   cugraph_error_t** error)
 {
 
-  *heterogeneous_fanout = reinterpret_cast<cugraph_sample_heterogeneous_fanout_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
+  *heterogeneous_fan_out = reinterpret_cast<cugraph_sample_heterogeneous_fanout_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
           new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_offsets)),
-          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fanout))});
+          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out))});
   
   return CUGRAPH_SUCCESS;
 }

From 8c1700949d7074bc80e1b004c1fc113d69841134 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 22 Aug 2024 17:44:09 -0700
Subject: [PATCH 024/117] rename variable for consistency

---
 cpp/include/cugraph_c/sampling_algorithms.h   | 10 +++++-----
 cpp/src/c_api/neighbor_sampling.cpp           | 19 +++++++++----------
 cpp/src/c_api/neighbor_sampling.hpp           |  2 +-
 .../_cugraph_c/sampling_algorithms.pxd        | 10 +++++-----
 .../pylibcugraph/neighbor_sample.pyx          | 10 +++++-----
 5 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index d117a329510..17dd53436f9 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -330,7 +330,7 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
 /// mimic
 typedef struct {
   int32_t align_;
-} cugraph_sample_heterogeneous_fanout_t;
+} cugraph_sample_heterogeneous_fan_out_t;
 
 /**
  * @brief       Create heterogeneous fanout
@@ -353,12 +353,12 @@ typedef struct {
  *                             be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_create_heterogeneous_fanout(
+cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
   const cugraph_type_erased_host_array_view_t* edge_type_offsets,
   const cugraph_type_erased_host_array_view_t* fanout,
-  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fanout,
   cugraph_error_t** error);
 
 /**
@@ -366,7 +366,7 @@ cugraph_error_code_t cugraph_create_heterogeneous_fanout(
  *
  * @param [in]    heterogeneous_fanout The edge type size and fanout values
  */
-void cugraph_heterogeneous_fanout_free(cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
+void cugraph_heterogeneous_fanout_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
 
 /**
  * @brief     Uniform Neighborhood Sampling
@@ -547,7 +547,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+  const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 019692829c4..46e6eacf938 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -77,7 +77,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
-  cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out_{nullptr};
+  cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const* heterogeneous_fan_out_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool is_biased_{false};
   bool do_expensive_check_{false};
@@ -94,7 +94,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
     cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
-    cugraph_sample_heterogeneous_fanout_t const* heterogeneous_fan_out,
+    cugraph_sample_heterogeneous_fan_out_t const* heterogeneous_fan_out,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool is_biased,
     bool do_expensive_check)
@@ -122,7 +122,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
 
       heterogeneous_fan_out_(
-        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(
+        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(
           heterogeneous_fan_out)),
   
       options_(options),
@@ -1007,7 +1007,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fan_out,
+  const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
@@ -1064,14 +1064,14 @@ cugraph_error_code_t cugraph_neighbor_sample(
   } else {
     
     CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)))
+                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
                    ->type_ == INT32,
                CUGRAPH_INVALID_INPUT,
                "edge type offsets type must be INT32",
                *error);
     
     CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fanout_t const*>(heterogeneous_fan_out)))
+                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
                    ->type_ == INT32,
                CUGRAPH_INVALID_INPUT,
                "fan_out values type must be INT32",
@@ -1186,16 +1186,15 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fanout(
+extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
   const cugraph_type_erased_host_array_view_t* edge_type_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  cugraph_sample_heterogeneous_fanout_t** heterogeneous_fan_out,
+  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
   cugraph_error_t** error)
 {
-
-  *heterogeneous_fan_out = reinterpret_cast<cugraph_sample_heterogeneous_fanout_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
+  *heterogeneous_fan_out = reinterpret_cast<cugraph_sample_heterogeneous_fan_out_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
           new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_offsets)),
           new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out))});
   
diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
index 9defb8a39dd..c4a2f8f3575 100644
--- a/cpp/src/c_api/neighbor_sampling.hpp
+++ b/cpp/src/c_api/neighbor_sampling.hpp
@@ -20,7 +20,7 @@
 namespace cugraph {
 namespace c_api {
 
-using cugraph_sample_heterogeneous_fanout_t = std::tuple<
+using cugraph_sample_heterogeneous_fan_out_t = std::tuple<
   cugraph_type_erased_host_array_t*, cugraph_type_erased_host_array_t*>;
 
 }  // namespace c_api
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 2074e88725f..0934507d56c 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -74,7 +74,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
         const cugraph_type_erased_host_array_view_t* fan_out,
-        const cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout,
+        const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
         const cugraph_sampling_options_t* options,
         bool_t is_biased,
         bool_t do_expensive_check,
@@ -123,19 +123,19 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
             cugraph_error_t** error
         )
 
-    ctypedef struct cugraph_sample_heterogeneous_fanout_t:
+    ctypedef struct cugraph_sample_heterogeneous_fan_out_t:
         pass
 
     cdef cugraph_error_code_t \
-        cugraph_create_heterogeneous_fanout(
+        cugraph_create_heterogeneous_fan_out(
             const cugraph_resource_handle_t* handle,
             cugraph_graph_t* graph,
             const cugraph_type_erased_host_array_view_t* edge_type_offsets,
             const cugraph_type_erased_host_array_view_t* fanout,
-            cugraph_sample_heterogeneous_fanout_t** heterogeneous_fanout,
+            cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fanout,
             cugraph_error_t** error
         )
 
     cdef void \
         cugraph_heterogeneous_fanout_free(
-            cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout);
+            cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
diff --git a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
index e8ce3e26992..8f40df523c7 100644
--- a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
@@ -56,8 +56,8 @@ from pylibcugraph._cugraph_c.algorithms cimport (
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
     cugraph_neighbor_sample,
-    cugraph_sample_heterogeneous_fanout_t,
-    cugraph_create_heterogeneous_fanout,
+    cugraph_sample_heterogeneous_fan_out_t,
+    cugraph_create_heterogeneous_fan_out,
     cugraph_heterogeneous_fanout_free,
 )
 from pylibcugraph.resource_handle cimport (
@@ -225,7 +225,7 @@ def neighbor_sample(ResourceHandle resource_handle,
     )
 
     cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
-    cdef cugraph_sample_heterogeneous_fanout_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fanout_t*>NULL
+    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
@@ -277,7 +277,7 @@ def neighbor_sample(ResourceHandle resource_handle,
                 len(h_fan_out[1]),
                 get_c_type_from_numpy_type(h_fan_out[1].dtype))
 
-        error_code = cugraph_create_heterogeneous_fanout(
+        error_code = cugraph_create_heterogeneous_fan_out(
             c_resource_handle_ptr,
             c_graph_ptr,
             fan_out_size_ptr,
@@ -286,7 +286,7 @@ def neighbor_sample(ResourceHandle resource_handle,
             &error_ptr
         )
 
-    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fanout")
+    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fan_out")
 
     cdef cugraph_sample_result_t* result_ptr
 

From 7b95c5e516b657692333cd1defbc758615d5f7bc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 29 Aug 2024 20:28:22 -0700
Subject: [PATCH 025/117] update neighbor sample API

---
 cpp/include/cugraph/sampling_functions.hpp    | 131 ++++-
 cpp/include/cugraph_c/sampling_algorithms.h   |  86 ++-
 cpp/src/c_api/neighbor_sampling.cpp           | 279 +++++++---
 cpp/src/sampling/neighbor_sampling_impl.hpp   |  65 ++-
 .../sampling/neighbor_sampling_mg_v32_e32.cpp |  73 ++-
 .../sampling/neighbor_sampling_mg_v32_e64.cpp |  64 ++-
 .../sampling/neighbor_sampling_mg_v64_e64.cpp |  64 ++-
 .../sampling/neighbor_sampling_sg_v32_e32.cpp |  70 ++-
 .../sampling/neighbor_sampling_sg_v32_e64.cpp |  64 ++-
 .../sampling/neighbor_sampling_sg_v64_e64.cpp |  64 ++-
 .../_cugraph_c/sampling_algorithms.pxd        |  23 +-
 .../heterogeneous_neighbor_sample.pyx         | 511 ++++++++++++++++++
 ...le.pyx => homogeneous_neighbor_sample.pyx} |  76 +--
 13 files changed, 1355 insertions(+), 215 deletions(-)
 create mode 100644 python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx
 rename python/pylibcugraph/pylibcugraph/{neighbor_sample.pyx => homogeneous_neighbor_sample.pyx} (90%)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 3140e9073a0..45114e06480 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -43,7 +43,8 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
 /**
  * @brief Uniform Neighborhood Sampling.
  * 
- * @deprecated  This API will be deleted, use neighbor_sample instead
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
+ * 'is_biased' set to false instead
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects from these outgoing neighbors to extract a subgraph.
@@ -142,7 +143,8 @@ uniform_neighbor_sample(
 /**
  * @brief Biased Neighborhood Sampling.
  * 
- * @deprecated  This API will be deleted, use neighbor_sample instead
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
+ * 'is_biased' set to true instead
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
@@ -244,14 +246,12 @@ biased_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
-
 /**
- * @brief Neighborhood Sampling.
+ * @brief Homogeneous Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
- * When branching out to select outgoing neighbors, either fan_out or heterogeneous_fan_out must
- * be provided but not both.
+ * The branching out to select outgoing neighbors is performed with homogeneous fanouts.
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
  * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
@@ -301,12 +301,10 @@ biased_neighbor_sample(
  * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
  * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
  * output rank.  The label span must be sorted in ascending order.
- * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level. When fan_out is provided, the sampling method uses the same fanout value for each type.
  * @param heterogeneous_fan_out Tuple of host spans defining branching out (fan-out) degree per
  * source vertex for each level in CSR style format. The first element of the tuple is the offset
- * array per edge type id and the second element correspond to the fanout values.
- * When heterogeneous_fan_out is provided, different fan_out values can be used for each edge type.
+ * array per edge type id and the second element corresponds to the fanout values.
+ * The sampling method can use different fan_out values for each edge type.
  * The fan-out offsets size must be proportional to the number of edge types and fan_out values.
  * @param return_hops boolean flag specifying if the hop information should be returned. 
  * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
@@ -321,7 +319,7 @@ biased_neighbor_sample(
  * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
  * optional int32_t hop, optional label_t label, optional size_t offsets)
  */
-// FIXME: Add flag for bias=True/False
+
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
@@ -338,7 +336,7 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<label_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
@@ -350,8 +348,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement                           = true,
@@ -359,6 +356,112 @@ neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
+/**
+ * @brief Homogeneous Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with homogeneous fanouts
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
+ * (optional) will be described below and is dependent upon the input parameters.
+ *
+ * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
+ * label and offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
+ * the label output has values.  This will also result in the output being sorted by vertex label.
+ * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
+ * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
+ *
+ * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
+ * label output has values.  This will also result in the output being sorted by vertex label.  The
+ * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
+ * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
+ * be shuffled so that all data with a particular label will be on the specified rank.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam label_t Type of label. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
+ * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
+ * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
+ * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
+ * uniform sampling.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * for the sampling.
+ * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
+ * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
+ * output rank.  The label span must be sorted in ascending order.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level. The sampling method uses the same fanout value for each type.
+ * @param return_hops boolean flag specifying if the hop information should be returned. 
+ * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
+ * DEFAULT)
+ * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
+ * multiple times with the same label, it will only be passed once (for each label) as a source
+ * for the next hop.  Default is false.
+ * @param with_replacement boolean flag specifying if random sampling is done with replacement
+ * (true); or, without replacement (false); default = true;
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement                           = true,
+  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
+  bool dedupe_sources                             = false,
+  bool do_expensive_check                         = false);
+
 /*
  * @brief renumber sampled edge list and compress to the (D)CSR|(D)CSC format.
  *
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 17dd53436f9..32f0b09ee6d 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -366,12 +366,13 @@ cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
  *
  * @param [in]    heterogeneous_fanout The edge type size and fanout values
  */
-void cugraph_heterogeneous_fanout_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
+void cugraph_heterogeneous_fan_out_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
 
 /**
  * @brief     Uniform Neighborhood Sampling
  * 
- * @deprecated  This API will be deleted, use cugraph_neighbor_sample instead
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
+ * 'is_biased' set to false instead
  *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
@@ -428,8 +429,9 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
 /**
  * @brief     Biased Neighborhood Sampling
  * 
- * @deprecated  This API will be deleted, use cugraph_neighbor_sample instead
- *
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
+ * 'is_biased' set to true instead
+ * 
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
  * of sampling requests in the same function call - which should improve GPU utilization.
@@ -487,9 +489,10 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_error_t** error);
 
 /**
- * @brief     Neighborhood Sampling
+ * @brief     Heterogeneous Neighborhood Sampling
  *
- * Returns a sample of the neighborhood around specified start vertices with edge biases or not.
+ * Returns a sample of the neighborhood around specified start vertices with edge biases or not
+ * and homogeneous fanout types.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -518,11 +521,9 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fanout values of type INT32
  * @param [in]  heterogeneous_fanout  Tuple of host arrays defining the fan out at each step in the
  * sampling algorithm. in CSR style format. The first element of the tuple is the offset array per
- * edge type id and the second element correspond to the fanout values.
+ * edge type id and the second element corresponds to the fanout values.
  * We only support type INT32 for both the offsets and the fanout values array.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
@@ -536,7 +537,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -546,7 +547,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
-  const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
@@ -554,6 +554,70 @@ cugraph_error_code_t cugraph_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
+/**
+ * @brief     Homogeneous Neighborhood Sampling
+ *
+ * Returns a sample of the neighborhood around specified start vertices with edge biases or not
+ * and homogeneous fanout types.
+ * Optionally, each start vertex can be associated with a label, allowing the caller to specify
+ * multiple batches of sampling requests in the same function call - which should improve GPU
+ * utilization.
+ *
+ * If label is NULL then all start vertices will be considered part of the same batch and the
+ * return value will not have a label column.
+ *
+ * @param [in]  handle       Handle for accessing resources
+ *  * @param [in,out] rng_state State of the random number generator, updated with each call
+ * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
+ *                           needs to be transposed
+ * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
+ * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
+ * @param [in]  start_vertices Device array of start vertices for the sampling
+ * @param [in]  start_vertex_labels  Device array of start vertex labels for the sampling.  The
+ * labels associated with each start vertex will be included in the output associated with results
+ * that were derived from that start vertex.  We only support label of type INT32. If label is
+ * NULL, the return data will not be labeled.
+ * @param [in]  label_list Device array of the labels included in @p start_vertex_labels.  If
+ * @p label_to_comm_rank is not specified this parameter is ignored.  If specified, label_list
+ * must be sorted in ascending order.
+ * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
+ * particular label should be shuffled in the output.  If not specifed the data is not organized in
+ * output.  If specified then the all data from @p label_list[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_labels is also specified
+ * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
+ * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
+ * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
+ *                           We only support fanout values of type INT32
+ * @param [in]  sampling_options
+ *                           Opaque pointer defining the sampling options.
+ * @param [in]  is_biased
+ *                           A flag specifying whether to run biased neighborhood sampling
+ *                           (if set to true) or uniform neighbor sampling.
+ * @param [in]  do_expensive_check
+ *                           A flag to run expensive checks for input arguments (if set to true)
+ * @param [out]  result      Output from the uniform_neighbor_sample call
+ * @param [out] error        Pointer to an error object storing details of any error.  Will
+ *                           be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_edge_property_view_t* edge_biases,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sampling_options_t* options,
+  bool_t is_biased,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error);
+
 /**
  * @deprecated This call should be replaced with cugraph_sample_result_get_majors
  * @brief     Get the source vertices from the sampling algorithm result
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 46e6eacf938..def04c021e7 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -218,45 +218,82 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_last(),
         do_expensive_check_);
       
+      rmm::device_uvector<vertex_t> src(0, handle_.get_stream());
+      rmm::device_uvector<vertex_t> dst(0, handle_.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt{std::nullopt};
+      std::optional<rmm::device_uvector<edge_t>> edge_id{std::nullopt};
+      std::optional<rmm::device_uvector<edge_type_t>> edge_type{std::nullopt};
+      std::optional<rmm::device_uvector<int32_t>> hop{std::nullopt};
+      std::optional<rmm::device_uvector<label_t>> edge_label{std::nullopt};
+      std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
       // FIXME: Consolidate 'fan_out_' and 'heterogeneous_fan_out_' into one
       // argument with std::variant
-      auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-        cugraph::neighbor_sample(
-          handle_,
-          rng_state_->rng_state_,
-          graph_view,
-          (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-          (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-          (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-          is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-          raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-          (start_vertex_labels_ != nullptr)
-            ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                   start_vertex_labels->size())
-            : std::nullopt,
-          (label_list_ != nullptr)
-            ? std::make_optional(std::make_tuple(
-                raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                 label_list_->size_},
-                raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                 label_to_comm_rank_->size_}))
-            : std::nullopt,
-          (fan_out_ != nullptr) ? std::make_optional<raft::host_span<const int>>(
-                                    fan_out_->as_type<const int>(), fan_out_->size_)
-                                : std::nullopt,
-          
-          (heterogeneous_fan_out_ != nullptr)
-            ? std::make_optional(std::make_tuple(
-                raft::host_span<const int>{std::get<0>(*heterogeneous_fan_out_)->as_type<int>(),
-                                           std::get<0>(*heterogeneous_fan_out_)->size_},
-                raft::host_span<const int>{std::get<1>(*heterogeneous_fan_out_)->as_type<int>(),
-                                           std::get<1>(*heterogeneous_fan_out_)->size_}))
-            : std::nullopt,
-          options_.return_hops_,
-          options_.with_replacement_,
-          options_.prior_sources_behavior_,
-          options_.dedupe_sources_,
-          do_expensive_check_);
+      //auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
+      if (heterogeneous_fan_out_ != nullptr) {
+        std::tie(src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets) =
+          cugraph::heterogeneous_neighbor_sample(
+            handle_,
+            rng_state_->rng_state_,
+            graph_view,
+            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+            (start_vertex_labels_ != nullptr)
+              ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                    start_vertex_labels->size())
+              : std::nullopt,
+            (label_list_ != nullptr)
+              ? std::make_optional(std::make_tuple(
+                  raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                  label_list_->size_},
+                  raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                  label_to_comm_rank_->size_}))
+              : std::nullopt,
+
+            std::make_tuple(
+              raft::host_span<const int>{std::get<0>(*heterogeneous_fan_out_)->as_type<int>(),
+                                        std::get<0>(*heterogeneous_fan_out_)->size_},
+              raft::host_span<const int>{std::get<1>(*heterogeneous_fan_out_)->as_type<int>(),
+                                        std::get<1>(*heterogeneous_fan_out_)->size_}),
+            options_.return_hops_,
+            options_.with_replacement_,
+            options_.prior_sources_behavior_,
+            options_.dedupe_sources_,
+            do_expensive_check_);
+      } else {
+
+        std::tie(src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets) =
+          cugraph::homogeneous_neighbor_sample(
+            handle_,
+            rng_state_->rng_state_,
+            graph_view,
+            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+            (start_vertex_labels_ != nullptr)
+              ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                    start_vertex_labels->size())
+              : std::nullopt,
+            (label_list_ != nullptr)
+              ? std::make_optional(std::make_tuple(
+                  raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                  label_list_->size_},
+                  raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                  label_to_comm_rank_->size_}))
+              : std::nullopt,
+            raft::host_span<const int>(
+                                      fan_out_->as_type<const int>(), fan_out_->size_),
+            options_.return_hops_,
+            options_.with_replacement_,
+            options_.prior_sources_behavior_,
+            options_.dedupe_sources_,
+            do_expensive_check_);
+
+      }
 
       std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
 
@@ -996,9 +1033,8 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-cugraph_error_code_t cugraph_neighbor_sample(
+cugraph_error_code_t cugraph_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
-  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
@@ -1007,6 +1043,88 @@ cugraph_error_code_t cugraph_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
+  cugraph_rng_state_t* rng_state,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error)
+{
+  auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
+
+  CAPI_EXPECTS(
+    (edge_biases != nullptr) ||
+      (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
+    CUGRAPH_INVALID_INPUT,
+    "edge_biases is required if the graph is not weighted",
+    *error);
+
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "must specify label_offsets if retain_seeds is true",
+               *error);
+
+  CAPI_EXPECTS((start_vertex_labels == nullptr) ||
+                 (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                    start_vertex_labels)
+                    ->type_ == INT32),
+               CUGRAPH_INVALID_INPUT,
+               "start_vertex_labels should be of type int",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_labels != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless start_vertex_labels is also specified",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless label_list is also specified",
+               *error);
+
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
+                 reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                   start_vertices)
+                   ->type_,
+               CUGRAPH_INVALID_INPUT,
+               "vertex type of graph and start_vertices must match",
+               *error);
+
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out should be of type int",
+    *error);
+
+  bool is_biased = true;
+
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    edge_biases,
+                                    start_vertices,
+                                    start_vertex_labels,
+                                    label_list,
+                                    label_to_comm_rank,
+                                    label_offsets,
+                                    fan_out,
+                                    nullptr,
+                                    std::move(options_cpp),
+                                    is_biased,
+                                    do_expensive_check};
+  return cugraph::c_api::run_algorithm(graph, functor, result, error);
+}
+
+cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_edge_property_view_t* edge_biases,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
@@ -1047,36 +1165,20 @@ cugraph_error_code_t cugraph_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "cannot specify label_to_comm_rank unless label_list is also specified",
                *error);
-  
-  CAPI_EXPECTS(!((fan_out != nullptr) && (heterogeneous_fan_out != nullptr)),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify both fan_out and heterogeneous_fan_out",
-               *error);
-  
-  if (fan_out != nullptr) {
-     CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                   fan_out)
-                   ->type_ == INT32,
-               CUGRAPH_INVALID_INPUT,
-               "fan_out type must be INT32",
-               *error);
-    
-  } else {
     
-    CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
-                   ->type_ == INT32,
-               CUGRAPH_INVALID_INPUT,
-               "edge type offsets type must be INT32",
-               *error);
-    
-    CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                   std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
-                   ->type_ == INT32,
-               CUGRAPH_INVALID_INPUT,
-               "fan_out values type must be INT32",
-               *error);
-  }
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                  std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
+                  ->type_ == INT32,
+              CUGRAPH_INVALID_INPUT,
+              "edge type offsets type must be INT32",
+              *error);
+  
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                  std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
+                  ->type_ == INT32,
+              CUGRAPH_INVALID_INPUT,
+              "fan_out values type must be INT32",
+              *error);
 
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -1086,7 +1188,6 @@ cugraph_error_code_t cugraph_neighbor_sample(
                "vertex type of graph and start_vertices must match",
                *error);
 
-
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
@@ -1096,7 +1197,7 @@ cugraph_error_code_t cugraph_neighbor_sample(
                                     label_list,
                                     label_to_comm_rank,
                                     label_offsets,
-                                    fan_out,
+                                    nullptr,
                                     heterogeneous_fan_out,
                                     std::move(options_cpp),
                                     is_biased,
@@ -1104,8 +1205,9 @@ cugraph_error_code_t cugraph_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-cugraph_error_code_t cugraph_biased_neighbor_sample(
+cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
@@ -1114,21 +1216,23 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  cugraph_rng_state_t* rng_state,
   const cugraph_sampling_options_t* options,
+  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
-  CAPI_EXPECTS(
-    (edge_biases != nullptr) ||
-      (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
-    CUGRAPH_INVALID_INPUT,
-    "edge_biases is required if the graph is not weighted",
-    *error);
-
+  if (is_biased) {
+    CAPI_EXPECTS(
+      (edge_biases != nullptr) ||
+        (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
+      CUGRAPH_INVALID_INPUT,
+      "edge_biases is required if the graph is not weighted",
+      *error);
+  }
+  
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify label_offsets if retain_seeds is true",
@@ -1151,7 +1255,14 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "cannot specify label_to_comm_rank unless label_list is also specified",
                *error);
-
+  
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                fan_out)
+                ->type_ == INT32,
+            CUGRAPH_INVALID_INPUT,
+            "fan_out type must be INT32",
+            *error);
+    
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                    start_vertices)
@@ -1160,14 +1271,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                "vertex type of graph and start_vertices must match",
                *error);
 
-  CAPI_EXPECTS(
-    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
-        ->type_ == INT32,
-    CUGRAPH_INVALID_INPUT,
-    "fan_out should be of type int",
-    *error);
-
-  bool is_biased = true;
 
   neighbor_sampling_functor functor{handle,
                                     rng_state,
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index fe9cceb1fd1..a49102ea7ad 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -472,6 +472,63 @@ biased_neighbor_sample(
     do_expensive_check);
 }
 
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          typename label_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<label_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check)
+{
+  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    edge_bias_view,
+    starting_vertices,
+    starting_vertex_labels,
+    label_to_output_comm_rank,
+    std::nullopt,
+    heterogeneous_fan_out,
+    return_hops,
+    with_replacement,
+    prior_sources_behavior,
+    dedupe_sources,
+    rng_state,
+    do_expensive_check);
+}
+
+
 
 
 
@@ -492,7 +549,7 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<label_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
@@ -504,9 +561,7 @@ 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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -524,7 +579,7 @@ neighbor_sample(
     starting_vertex_labels,
     label_to_output_comm_rank,
     fan_out,
-    heterogeneous_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 5aa325c8578..917e8fcc6e1 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
@@ -127,15 +127,71 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<float>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<double>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<float>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
@@ -147,9 +203,7 @@ 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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -164,7 +218,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
@@ -176,13 +230,14 @@ 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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
   bool dedupe_sources,
   bool do_expensive_check);
 
+
+
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
index 913e7d8a62e..71661090023 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
@@ -136,7 +136,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
@@ -148,8 +148,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -165,7 +164,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
@@ -177,8 +176,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -186,4 +184,58 @@ neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
index 77ab2e476ac..2a1c0d43ab3 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
@@ -135,7 +135,7 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
@@ -147,8 +147,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -164,7 +163,7 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
@@ -176,8 +175,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -185,4 +183,58 @@ neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
index 05ccdfca1f7..c58e1659d7c 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
@@ -127,15 +127,71 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<float>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<double>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<float>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
@@ -147,9 +203,7 @@ 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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
@@ -164,7 +218,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
@@ -176,9 +230,7 @@ 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,
   bool return_hops,
   bool with_replacement,
   prior_sources_behavior_t prior_sources_behavior,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
index 178b1322d89..2a70085d490 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
@@ -135,7 +135,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
@@ -147,8 +147,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -164,7 +163,7 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
@@ -176,8 +175,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -185,4 +183,58 @@ neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
index ba1edd3f630..6d045ad092d 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
@@ -135,7 +135,7 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
@@ -147,8 +147,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -164,7 +163,7 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
@@ -176,8 +175,7 @@ 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>>>
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
     heterogeneous_fan_out,
   bool return_hops,
   bool with_replacement,
@@ -185,4 +183,58 @@ neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  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,
+  raft::host_span<int32_t const> fan_out,
+  bool return_hops,
+  bool with_replacement,
+  prior_sources_behavior_t prior_sources_behavior,
+  bool dedupe_sources,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 0934507d56c..6862ae6054c 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -63,7 +63,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_error_t** error
     )
 
-    cdef cugraph_error_code_t cugraph_neighbor_sample(
+    cdef cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
         const cugraph_resource_handle_t* handle,
         cugraph_rng_state_t* rng_state,
         cugraph_graph_t* graph,
@@ -73,7 +73,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_list,
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
-        const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
         const cugraph_sampling_options_t* options,
         bool_t is_biased,
@@ -82,6 +81,24 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_error_t** error
     )
 
+    cdef cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
+        const cugraph_resource_handle_t* handle,
+        cugraph_rng_state_t* rng_state,
+        cugraph_graph_t* graph,
+        const cugraph_edge_property_view_t* edge_biases,
+        const cugraph_type_erased_device_array_view_t* start_vertices,
+        const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+        const cugraph_type_erased_device_array_view_t* label_list,
+        const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+        const cugraph_type_erased_device_array_view_t* label_offsets,
+        const cugraph_type_erased_host_array_view_t* fan_out,
+        const cugraph_sampling_options_t* options,
+        bool_t is_biased,
+        bool_t do_expensive_check,
+        cugraph_sample_result_t** result,
+        cugraph_error_t** error
+    )
+
     cdef cugraph_error_code_t cugraph_biased_neighbor_sample(
         const cugraph_resource_handle_t* handle,
         cugraph_graph_t* graph,
@@ -137,5 +154,5 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         )
 
     cdef void \
-        cugraph_heterogeneous_fanout_free(
+        cugraph_heterogeneous_fan_out_free(
             cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx
new file mode 100644
index 00000000000..0f999adf93d
--- /dev/null
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx
@@ -0,0 +1,511 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    bool_t,
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.properties cimport (
+    cugraph_edge_property_view_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.algorithms cimport (
+    cugraph_sample_result_t,
+    cugraph_prior_sources_behavior_t,
+    cugraph_compression_type_t,
+    cugraph_sampling_options_t,
+    cugraph_sampling_options_create,
+    cugraph_sampling_options_free,
+    cugraph_sampling_set_with_replacement,
+    cugraph_sampling_set_return_hops,
+    cugraph_sampling_set_prior_sources_behavior,
+    cugraph_sampling_set_dedupe_sources,
+    cugraph_sampling_set_renumber_results,
+    cugraph_sampling_set_compress_per_hop,
+    cugraph_sampling_set_compression_type,
+    cugraph_sampling_set_retain_seeds,
+)
+from pylibcugraph._cugraph_c.sampling_algorithms cimport (
+    cugraph_heterogeneous_neighbor_sample,
+    cugraph_sample_heterogeneous_fan_out_t,
+    cugraph_create_heterogeneous_fan_out,
+    cugraph_heterogeneous_fan_out_free,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+)
+from pylibcugraph.internal_types.sampling_result cimport (
+    SamplingResult,
+)
+from pylibcugraph._cugraph_c.random cimport (
+    cugraph_rng_state_t
+)
+from pylibcugraph.random cimport (
+    CuGraphRandomState
+)
+import warnings
+
+# TODO accept cupy/numpy random state in addition to raw seed.
+def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
+                    _GPUGraph input_graph,
+                    start_list,
+                    h_fan_out,
+                    bool_t is_biased,
+                    bool_t with_replacement,
+                    bool_t do_expensive_check,
+                    with_edge_properties=False,
+                    batch_id_list=None,
+                    label_list=None,
+                    label_to_output_comm_rank=None,
+                    label_offsets=None,
+                    prior_sources_behavior=None,
+                    deduplicate_sources=False,
+                    return_hops=False,
+                    renumber=False,
+                    retain_seeds=False,
+                    compression='COO',
+                    compress_per_hop=False,
+                    random_state=None,
+                    return_dict=False,):
+    """
+    # FIXME: Deprecate uniform_neighbor_sample
+    Does both uniform and biased neighborhood sampling, which samples nodes from
+    a graph based on the current node's neighbors, with a corresponding fan_out
+    value at each hop.
+
+    # FIXME: biased neighbor sampling is not yet exposed to PLC.
+
+    Parameters
+    ----------
+    resource_handle: ResourceHandle
+        Handle to the underlying device and host resources needed for
+        referencing data and running algorithms.
+
+    input_graph : SGGraph or MGGraph
+        The input graph, for either Single or Multi-GPU operations.
+
+    start_list: device array type
+        Device array containing the list of starting vertices for sampling.
+
+    edge_biases: FIXE: update this - Create edge_biases of type
+    'cugraph_edge_property_view_t' - edge biases not yet supported.
+
+    h_fan_out: tuple of numpy array type
+        Device array containing the branching out (fan-out) degrees per
+        starting vertex for each hop level in CSR style format. The first
+        element of the tuple is the offset array per edge type id and the second
+        element corresponds to the fan_out values.
+        The sampling method can use different fan_out values for each edge type.
+    
+    is_biased: bool
+        If false, sampling procedure is done uniform otherwise with biases.
+
+    with_replacement: bool
+        If true, sampling procedure is done with replacement (the same vertex
+        can be selected multiple times in the same step).
+
+    do_expensive_check: bool
+        If True, performs more extensive tests on the inputs to ensure
+        validitity, at the expense of increased run time.
+
+    with_edge_properties: bool
+        If True, returns the edge properties of each edges along with the
+        edges themselves.  Will result in an error if the provided graph
+        does not have edge properties.
+
+    batch_id_list: list[int32] (Optional)
+        List of int32 batch ids that is returned with each edge.  Optional
+        argument, defaults to NULL, returning nothing.
+
+    label_list: list[int32] (Optional)
+        List of unique int32 batch ids.  Required if also passing the
+        label_to_output_comm_rank flag.  Default to NULL (does nothing)
+
+    label_to_output_comm_rank: list[int32] (Optional)
+        Maps the unique batch ids in label_list to the rank of the
+        worker that should hold results for that batch id.
+        Defaults to NULL (does nothing)
+
+    label_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    prior_sources_behavior: str (Optional)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+
+    deduplicate_sources: bool (Optional)
+        If True, will deduplicate the source list before sampling.
+        Defaults to False.
+
+    renumber: bool (Optional)
+        If True, will renumber the sources and destinations on a
+        per-batch basis and return the renumber map and batch offsets
+        in additional to the standard returns.
+
+    retain_seeds: bool (Optional)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+        Defaults to False.
+
+    compression: str (Optional)
+        Options: COO (default), CSR, CSC, DCSR, DCSR
+        Sets the compression format for the returned samples.
+
+    compress_per_hop: bool (Optional)
+        If False (default), will create a compressed edgelist for the
+        entire batch.
+        If True, will create a separate compressed edgelist per hop within
+        a batch.
+
+    random_state: int (Optional)
+        Random state to use when generating samples.  Optional argument,
+        defaults to a hash of process id, time, and hostname.
+        (See pylibcugraph.random.CuGraphRandomState)
+
+    return_dict: bool (Optional)
+        Whether to return a dictionary instead of a tuple.
+        Optional argument, defaults to False, returning a tuple.
+        This argument will eventually be deprecated in favor
+        of always returning a dictionary.
+
+    Returns
+    -------
+    A tuple of device arrays, where the first and second items in the tuple
+    are device arrays containing the starting and ending vertices of each
+    walk respectively, the third item in the tuple is a device array
+    containing the start labels, and the fourth item in the tuple is a device
+    array containing the indices for reconstructing paths.
+
+    If renumber was set to True, then the fifth item in the tuple is a device
+    array containing the renumber map, and the sixth item in the tuple is a
+    device array containing the renumber map offsets (which delineate where
+    the renumber map for each batch starts).
+
+    """
+    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
+        resource_handle.c_resource_handle_ptr
+    )
+
+    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_offsets_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+
+    cdef bool_t c_deduplicate_sources = deduplicate_sources
+    cdef bool_t c_return_hops = return_hops
+    cdef bool_t c_renumber = renumber
+    cdef bool_t c_compress_per_hop = compress_per_hop
+
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+    cdef uintptr_t ai_fan_out_offsets_ptr
+    cdef uintptr_t ai_fan_out_values_ptr
+
+
+    assert_CAI_type(start_list, "start_list")
+    assert_CAI_type(batch_id_list, "batch_id_list", True)
+    assert_CAI_type(label_list, "label_list", True)
+    assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
+    assert_CAI_type(label_offsets, "label_offsets", True)
+
+    assert_AI_type(h_fan_out[0], "h_fan_out_size")
+    assert_AI_type(h_fan_out[1], "h_fan_out_values")
+    ai_fan_out_offsets_ptr = \
+        h_fan_out[0].__array_interface__["data"][0]
+    ai_fan_out_values_ptr = \
+        h_fan_out[1].__array_interface__["data"][0]
+
+    fan_out_offsets_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_offsets_ptr,
+            len(h_fan_out[0]),
+            get_c_type_from_numpy_type(h_fan_out[0].dtype))
+
+    fan_out_values_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_values_ptr,
+            len(h_fan_out[1]),
+            get_c_type_from_numpy_type(h_fan_out[1].dtype))
+
+    error_code = cugraph_create_heterogeneous_fan_out(
+        c_resource_handle_ptr,
+        c_graph_ptr,
+        fan_out_offsets_ptr,
+        fan_out_values_ptr,
+        &heterogeneous_fan_out_ptr,
+        &error_ptr
+    )
+
+    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fan_out")
+
+    cdef cugraph_sample_result_t* result_ptr
+
+    cdef uintptr_t cai_start_ptr = \
+        start_list.__cuda_array_interface__["data"][0]
+
+    cdef uintptr_t cai_batch_id_ptr
+    if batch_id_list is not None:
+        cai_batch_id_ptr = \
+            batch_id_list.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_list_ptr
+    if label_list is not None:
+        cai_label_list_ptr = \
+            label_list.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_to_output_comm_rank_ptr
+    if label_to_output_comm_rank is not None:
+        cai_label_to_output_comm_rank_ptr = \
+            label_to_output_comm_rank.__cuda_array_interface__['data'][0]
+
+    cdef uintptr_t cai_label_offsets_ptr
+    if label_offsets is not None:
+        cai_label_offsets_ptr = \
+            label_offsets.__cuda_array_interface__['data'][0]
+
+    cdef cugraph_type_erased_device_array_view_t* start_ptr = \
+        cugraph_type_erased_device_array_view_create(
+            <void*>cai_start_ptr,
+            len(start_list),
+            get_c_type_from_numpy_type(start_list.dtype))
+
+    cdef cugraph_type_erased_device_array_view_t* batch_id_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if batch_id_list is not None:
+        batch_id_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_batch_id_ptr,
+                len(batch_id_list),
+                get_c_type_from_numpy_type(batch_id_list.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_list_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if label_list is not None:
+        label_list_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_list_ptr,
+                len(label_list),
+                get_c_type_from_numpy_type(label_list.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_to_output_comm_rank_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if label_to_output_comm_rank is not None:
+        label_to_output_comm_rank_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_to_output_comm_rank_ptr,
+                len(label_to_output_comm_rank),
+                get_c_type_from_numpy_type(label_to_output_comm_rank.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if retain_seeds:
+        if label_offsets is None:
+            raise ValueError("Must provide label offsets if retain_seeds is True")
+        label_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_label_offsets_ptr,
+                len(label_offsets),
+                get_c_type_from_numpy_type(label_offsets.dtype)
+            )
+
+    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
+
+    cdef cugraph_rng_state_t* rng_state_ptr = \
+        cg_rng_state.rng_state_ptr
+
+    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
+    if prior_sources_behavior is None:
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
+    elif prior_sources_behavior == 'carryover':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
+    elif prior_sources_behavior == 'exclude':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
+    else:
+        raise ValueError(
+            f'Invalid option {prior_sources_behavior}'
+            ' for prior sources behavior'
+        )
+
+    cdef cugraph_compression_type_t compression_behavior_e
+    if compression is None or compression == 'COO':
+        compression_behavior_e = cugraph_compression_type_t.COO
+    elif compression == 'CSR':
+        compression_behavior_e = cugraph_compression_type_t.CSR
+    elif compression == 'CSC':
+        compression_behavior_e = cugraph_compression_type_t.CSC
+    elif compression == 'DCSR':
+        compression_behavior_e = cugraph_compression_type_t.DCSR
+    elif compression == 'DCSC':
+        compression_behavior_e = cugraph_compression_type_t.DCSC
+    else:
+        raise ValueError(
+            f'Invalid option {compression}'
+            ' for compression type'
+        )
+
+    cdef cugraph_sampling_options_t* sampling_options
+    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
+
+    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
+    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
+    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
+    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
+    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
+    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
+    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
+    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
+
+    error_code = cugraph_heterogeneous_neighbor_sample(
+        c_resource_handle_ptr,
+        rng_state_ptr,
+        c_graph_ptr,
+        <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
+        start_ptr,
+        batch_id_ptr,
+        label_list_ptr,
+        label_to_output_comm_rank_ptr,
+        label_offsets_ptr,
+        heterogeneous_fan_out_ptr,
+        sampling_options,
+        is_biased,
+        do_expensive_check,
+        &result_ptr,
+        &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_heterogeneous_neighbor_sample")
+
+    # Free the sampling options
+    cugraph_sampling_options_free(sampling_options)
+
+    if isinstance(h_fan_out, tuple):
+        cugraph_heterogeneous_fan_out_free(heterogeneous_fan_out_ptr)
+
+    # Free the two input arrays that are no longer needed.
+    cugraph_type_erased_device_array_view_free(start_ptr)
+    cugraph_type_erased_host_array_view_free(fan_out_ptr)
+    if batch_id_list is not None:
+        cugraph_type_erased_device_array_view_free(batch_id_ptr)
+    if label_offsets is not None:
+        cugraph_type_erased_device_array_view_free(label_offsets_ptr)
+
+    # Have the SamplingResult instance assume ownership of the result data.
+    result = SamplingResult()
+    result.set_ptr(result_ptr)
+
+    # Get cupy "views" of the individual arrays to return. These each increment
+    # the refcount on the SamplingResult instance which will keep the data alive
+    # until all references are removed and the GC runs.
+    # TODO Return everything that isn't null in release 23.12
+    if with_edge_properties:
+        cupy_majors = result.get_majors()
+        cupy_major_offsets = result.get_major_offsets()
+        cupy_minors = result.get_minors()
+        cupy_edge_weights = result.get_edge_weights()
+        cupy_edge_ids = result.get_edge_ids()
+        cupy_edge_types = result.get_edge_types()
+        cupy_batch_ids = result.get_batch_ids()
+        cupy_label_hop_offsets = result.get_label_hop_offsets()
+
+        if renumber:
+            cupy_renumber_map = result.get_renumber_map()
+            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+            # TODO drop the placeholder for hop ids in release 23.12
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': None,
+                    'renumber_map': cupy_renumber_map,
+                    'renumber_map_offsets': cupy_renumber_map_offsets
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+        else:
+            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': cupy_hop_ids,
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
+
+    else:
+        # TODO this is deprecated, remove it in release 23.12
+        warnings.warn(
+            "Calling uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
+            "with_edge_properties=True option, returning whatever properties are in the graph.",
+            FutureWarning,
+        )
+
+        cupy_sources = result.get_sources()
+        cupy_destinations = result.get_destinations()
+        cupy_indices = result.get_indices()
+
+        if return_dict:
+            return {
+                'sources': cupy_sources,
+                'destinations': cupy_destinations,
+                'indices': cupy_indices
+            }
+        else:
+            return (cupy_sources, cupy_destinations, cupy_indices)
diff --git a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx
similarity index 90%
rename from python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
rename to python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx
index 8f40df523c7..6b89980f171 100644
--- a/python/pylibcugraph/pylibcugraph/neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx
@@ -55,10 +55,9 @@ from pylibcugraph._cugraph_c.algorithms cimport (
     cugraph_sampling_set_retain_seeds,
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
-    cugraph_neighbor_sample,
+    cugraph_homogeneous_neighbor_sample,
     cugraph_sample_heterogeneous_fan_out_t,
-    cugraph_create_heterogeneous_fan_out,
-    cugraph_heterogeneous_fanout_free,
+    cugraph_heterogeneous_fan_out_free,
 )
 from pylibcugraph.resource_handle cimport (
     ResourceHandle,
@@ -84,7 +83,7 @@ from pylibcugraph.random cimport (
 import warnings
 
 # TODO accept cupy/numpy random state in addition to raw seed.
-def neighbor_sample(ResourceHandle resource_handle,
+def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
                     _GPUGraph input_graph,
                     start_list,
                     h_fan_out,
@@ -108,7 +107,7 @@ def neighbor_sample(ResourceHandle resource_handle,
     """
     # FIXME: Deprecate uniform_neighbor_sample
     Does both uniform and biased neighborhood sampling, which samples nodes from
-    a graph based on the current node's neighbors, with a corresponding fanout
+    a graph based on the current node's neighbors, with a corresponding fan_out
     value at each hop.
 
     # FIXME: biased neighbor sampling is not yet exposed to PLC.
@@ -128,11 +127,13 @@ def neighbor_sample(ResourceHandle resource_handle,
     edge_biases: FIXE: update this - Create edge_biases of type
     'cugraph_edge_property_view_t' - edge biases not yet supported.
 
-    h_fan_out: numpy array type or tuple of numpy array type
-        Device array containing the brancing out (fan-out) degrees per
-        starting vertex for each hop level.
+    h_fan_out: numpy array type
+        Device array containing the branching out (fan-out) degrees per
+        starting vertex for each hop level. The sampling method uses the same
+        fan_out value for each type.
     
-    is_biased: bool # FIXME: Update docstrings
+    is_biased: bool
+        If false, sampling procedure is done uniform otherwise with biases.
 
     with_replacement: bool
         If true, sampling procedure is done with replacement (the same vertex
@@ -225,7 +226,7 @@ def neighbor_sample(ResourceHandle resource_handle,
     )
 
     cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
-    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
+    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
@@ -247,44 +248,16 @@ def neighbor_sample(ResourceHandle resource_handle,
     assert_CAI_type(label_list, "label_list", True)
     assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
     assert_CAI_type(label_offsets, "label_offsets", True)
-    if not isinstance(h_fan_out, tuple):
-        assert_AI_type(h_fan_out, "h_fan_out")
-        ai_fan_out_ptr = \
-            h_fan_out.__array_interface__["data"][0]
-
-        fan_out_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_ptr,
-                len(h_fan_out),
-                get_c_type_from_numpy_type(h_fan_out.dtype))
-    else:
-        assert_AI_type(h_fan_out[0], "h_fan_out_size")
-        assert_AI_type(h_fan_out[1], "h_fan_out_values")
-        ai_fan_out_size_ptr = \
-            h_fan_out[0].__array_interface__["data"][0]
-        ai_fan_out_values_ptr = \
-            h_fan_out[1].__array_interface__["data"][0]
-
-        fan_out_size_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_size_ptr,
-                len(h_fan_out[0]),
-                get_c_type_from_numpy_type(h_fan_out[0].dtype))
-
-        fan_out_values_ptr = \
-            cugraph_type_erased_host_array_view_create(
-                <void*>ai_fan_out_values_ptr,
-                len(h_fan_out[1]),
-                get_c_type_from_numpy_type(h_fan_out[1].dtype))
-
-        error_code = cugraph_create_heterogeneous_fan_out(
-            c_resource_handle_ptr,
-            c_graph_ptr,
-            fan_out_size_ptr,
-            fan_out_values_ptr,
-            &heterogeneous_fanout_ptr,
-            &error_ptr
-        )
+
+    assert_AI_type(h_fan_out, "h_fan_out")
+    ai_fan_out_ptr = \
+        h_fan_out.__array_interface__["data"][0]
+
+    fan_out_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_ptr,
+            len(h_fan_out),
+            get_c_type_from_numpy_type(h_fan_out.dtype))
 
     assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fan_out")
 
@@ -405,7 +378,7 @@ def neighbor_sample(ResourceHandle resource_handle,
     cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
     cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
 
-    error_code = cugraph_neighbor_sample(
+    error_code = cugraph_homogeneous_neighbor_sample(
         c_resource_handle_ptr,
         rng_state_ptr,
         c_graph_ptr,
@@ -416,19 +389,18 @@ def neighbor_sample(ResourceHandle resource_handle,
         label_to_output_comm_rank_ptr,
         label_offsets_ptr,
         fan_out_ptr,
-        heterogeneous_fanout_ptr,
         sampling_options,
         is_biased,
         do_expensive_check,
         &result_ptr,
         &error_ptr)
-    assert_success(error_code, error_ptr, "cugraph_neighbor_sample")
+    assert_success(error_code, error_ptr, "cugraph_homogeneous_neighbor_sample")
 
     # Free the sampling options
     cugraph_sampling_options_free(sampling_options)
 
     if isinstance(h_fan_out, tuple):
-        cugraph_heterogeneous_fanout_free(heterogeneous_fanout_ptr)
+        cugraph_heterogeneous_fan_out_free(fan_out_ptr)
 
     # Free the two input arrays that are no longer needed.
     cugraph_type_erased_device_array_view_free(start_ptr)

From 19fc7655fa644db96fe1f69620c3576f4178b517 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 30 Aug 2024 08:15:45 -0700
Subject: [PATCH 026/117] remove outdated code

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index a49102ea7ad..3bf569956d8 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -69,13 +69,6 @@ neighbor_sample_impl(
   raft::random::RngState& rng_state,
   bool do_expensive_check)
 {
-#ifdef NO_CUGRAPH_OPS  // FIXME: this is relevant only when edge_bias_view.has_value() is false,
-                       // this ifdef statement will be removed once we migrate relevant cugraph-ops
-                       // functions to cugraph
-  CUGRAPH_FAIL(
-    "neighbor_sample_impl not supported in this configuration, built with NO_CUGRAPH_OPS");
-#else
-  static_assert(std::is_floating_point_v<bias_t>);
 
   if (fan_out) {
     CUGRAPH_EXPECTS((*fan_out).size() > 0,

From e30766c76218b66fceff6bfd2ac960f93a5d11dd Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 30 Aug 2024 08:17:29 -0700
Subject: [PATCH 027/117] remove outdated comment

---
 cpp/include/cugraph_c/sampling_algorithms.h | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 32f0b09ee6d..a98384c3016 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -322,12 +322,6 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
 /**
  * @brief     Opaque neighborhood sampling heterogeneous fanout type
  */
-// FIXME: internal representation should be tuple instead of pairs - Make it more generic (tuple)
-// cugraph_device_tuple_t, host_device_tuple_t,
-// dictionary, key and array
-// translate dictionary to a tuple. Add to the draft PR the PLC layer.
-// Concatenate to build the 3 arrays from the PLC layer
-/// mimic
 typedef struct {
   int32_t align_;
 } cugraph_sample_heterogeneous_fan_out_t;

From 5dd66f2a66abbb225bd0a186937505b85bbd9151 Mon Sep 17 00:00:00 2001
From: Charles Hastings <charlesh@nvidia.com>
Date: Wed, 4 Sep 2024 14:57:53 -0700
Subject: [PATCH 028/117] first cut at new sampling function definition to
 clean up things before MTMG integration

---
 cpp/include/cugraph/sampling_functions.hpp | 224 +++++++++++++++++++++
 1 file changed, 224 insertions(+)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 783cd3a7e2b..1799209242a 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -240,6 +240,230 @@ biased_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
+/**
+ * @brief Defines flags for sampling
+ *
+ * Note that this uses the Builder Pattern to allow specification of options.
+ * An example of how to build a complex option:
+ *
+ * sampling_flags_t x = sampling_flags_t{}.set_dedupe_sources(true);
+ *
+ */
+struct sampling_flags_t {
+  /**
+   * Specifies how to handle prior sources. Default is DEFAULT.
+   */
+  prior_sources_behavior_t prior_sources_behavior{};
+
+  /**
+   * Specifies if the hop information should be returned.  Default is false.
+   */
+  bool return_hops{false};
+
+  /**
+   * If true then if a vertex v appears as a destination in hop X multiple times
+   * with the same label, it will only be passed once (for each label) as a source
+   * for the next hop.  Default is false.
+   */
+  bool dedupe_sources{false};
+
+  /**
+   * Specifies if random sampling is done with replacement
+   *   (true) or without replacement (false).  Default is true.
+   */
+  bool with_replacement{true};
+
+  /**
+   * Set prior_sources_behavior to the specified value
+   */
+  sampling_flags_t set_prior_sources_behavior(prior_sources_behavior_t value)
+  {
+    return sampling_flags_t{value, return_hops, dedupe_sources, with_replacement};
+  }
+
+  /**
+   * Set return_hops to the specified value
+   */
+  sampling_flags_t set_return_hops(bool value)
+  {
+    return sampling_flags_t{prior_sources_behavior, value, dedupe_sources, with_replacement};
+  }
+
+  /**
+   * Set dedupe_sources to the specified value
+   */
+  sampling_flags_t set_dedupe_sources(bool value)
+  {
+    return sampling_flags_t{prior_sources_behavior, return_hops, value, with_replacement};
+  }
+
+  /**
+   * Set with_replacement to the specified value
+   */
+  sampling_flags_t set_with_replacement(bool value)
+  {
+    return sampling_flags_t{prior_sources_behavior, return_hops, dedupe_sources, value};
+  }
+};
+
+/**
+ * @brief Uniform Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The offsets array (optional) identifies the offset for each label.
+ *
+ * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * in the return will be a CSR-style offsets array to identify the beginning of each label range in
+ * the output vectors.
+ *
+ * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
+ * for a particular label are returned on the specified rank.  This will result in the offsets array
+ * on other GPUs indicating that there are no entries for that label (`offsets[i] == offsets[i+1]`).
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_offsets Optional device span of offsets identifying the range of
+ * starting vertex values for this label.
+ * @param label_to_output_comm_rank Optional device span identifying which rank should get each
+ * vertex label.  This should be the same on each rank.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level
+ * @param flags A set of flags indicating which sampling features should be used.l
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional size_t offsets)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check = false);
+
+/**
+ * @brief Biased Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The offsets array (optional) identifies the offset for each label.
+ *
+ * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * in the return will be a CSR-style offsets array to identify the beginning of each label range in
+ * the output vectors.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
+ * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
+ * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
+ * corresponding edge can never be selected.
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_offsets Optional device span of offsets identifying the range of
+ * starting vertex values for this label.
+ * @param label_to_output_comm_rank Optional device span identifying which rank should get each
+ * vertex label.  This should be the same on each rank.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level
+ * @param flags A set of flags indicating which sampling features should be used.l
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional size_t offsets)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check = false);
+
 /*
  * @brief renumber sampled edge list and compress to the (D)CSR|(D)CSC format.
  *

From 4b2764cc2046301451c9974feed19bc3d9520bae Mon Sep 17 00:00:00 2001
From: Charles Hastings <charlesh@nvidia.com>
Date: Thu, 5 Sep 2024 14:39:16 -0700
Subject: [PATCH 029/117] updates to remove builder pattern, also rename
 functions and mark old ones as deprecated

---
 cpp/include/cugraph/sampling_functions.hpp | 46 +++-------------------
 1 file changed, 6 insertions(+), 40 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 1799209242a..451094e1b48 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -43,6 +43,8 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
 /**
  * @brief Uniform Neighborhood Sampling.
  *
+ * @deprecated Replaced with homogeneous_uniform_neighbor_sample
+ *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects from these outgoing neighbors to extract a subgraph.
  *
@@ -140,6 +142,8 @@ uniform_neighbor_sample(
 /**
  * @brief Biased Neighborhood Sampling.
  *
+ * @deprecated Replaced with homogeneous_biased_neighbor_sample
+ *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
  *
@@ -242,12 +246,6 @@ biased_neighbor_sample(
 
 /**
  * @brief Defines flags for sampling
- *
- * Note that this uses the Builder Pattern to allow specification of options.
- * An example of how to build a complex option:
- *
- * sampling_flags_t x = sampling_flags_t{}.set_dedupe_sources(true);
- *
  */
 struct sampling_flags_t {
   /**
@@ -272,38 +270,6 @@ struct sampling_flags_t {
    *   (true) or without replacement (false).  Default is true.
    */
   bool with_replacement{true};
-
-  /**
-   * Set prior_sources_behavior to the specified value
-   */
-  sampling_flags_t set_prior_sources_behavior(prior_sources_behavior_t value)
-  {
-    return sampling_flags_t{value, return_hops, dedupe_sources, with_replacement};
-  }
-
-  /**
-   * Set return_hops to the specified value
-   */
-  sampling_flags_t set_return_hops(bool value)
-  {
-    return sampling_flags_t{prior_sources_behavior, value, dedupe_sources, with_replacement};
-  }
-
-  /**
-   * Set dedupe_sources to the specified value
-   */
-  sampling_flags_t set_dedupe_sources(bool value)
-  {
-    return sampling_flags_t{prior_sources_behavior, return_hops, value, with_replacement};
-  }
-
-  /**
-   * Set with_replacement to the specified value
-   */
-  sampling_flags_t set_with_replacement(bool value)
-  {
-    return sampling_flags_t{prior_sources_behavior, return_hops, dedupe_sources, value};
-  }
 };
 
 /**
@@ -370,7 +336,7 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-uniform_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
@@ -449,7 +415,7 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-biased_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,

From 4c1c61016cc9a5a06608cfab55df049a970eab16 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 9 Sep 2024 08:33:36 -0700
Subject: [PATCH 030/117] add implementation of heterogeneous neighborhood
 sampling

---
 cpp/CMakeLists.txt                            |  12 +-
 cpp/src/community/k_truss_impl.cuh            |  19 +-
 cpp/src/sampling/neighbor_sampling_impl.hpp   | 405 ++++++++++++++++--
 ...32.cpp => neighbor_sampling_mg_v32_e32.cu} |   0
 ...64.cpp => neighbor_sampling_mg_v32_e64.cu} |   0
 ...64.cpp => neighbor_sampling_mg_v64_e64.cu} |   0
 ...32.cpp => neighbor_sampling_sg_v32_e32.cu} |   0
 ...64.cpp => neighbor_sampling_sg_v32_e64.cu} |   0
 ...64.cpp => neighbor_sampling_sg_v64_e64.cu} |   0
 cpp/tests/community/k_truss_test.cpp          |  54 +--
 10 files changed, 391 insertions(+), 99 deletions(-)
 rename cpp/src/sampling/{neighbor_sampling_mg_v32_e32.cpp => neighbor_sampling_mg_v32_e32.cu} (100%)
 rename cpp/src/sampling/{neighbor_sampling_mg_v32_e64.cpp => neighbor_sampling_mg_v32_e64.cu} (100%)
 rename cpp/src/sampling/{neighbor_sampling_mg_v64_e64.cpp => neighbor_sampling_mg_v64_e64.cu} (100%)
 rename cpp/src/sampling/{neighbor_sampling_sg_v32_e32.cpp => neighbor_sampling_sg_v32_e32.cu} (100%)
 rename cpp/src/sampling/{neighbor_sampling_sg_v32_e64.cpp => neighbor_sampling_sg_v32_e64.cu} (100%)
 rename cpp/src/sampling/{neighbor_sampling_sg_v64_e64.cpp => neighbor_sampling_sg_v64_e64.cu} (100%)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 441627cabce..c37d0e8f35a 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -326,12 +326,12 @@ set(CUGRAPH_SOURCES
     src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu
     src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu
     src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
-    src/sampling/neighbor_sampling_mg_v32_e64.cpp
-    src/sampling/neighbor_sampling_mg_v32_e32.cpp
-    src/sampling/neighbor_sampling_mg_v64_e64.cpp
-    src/sampling/neighbor_sampling_sg_v32_e64.cpp
-    src/sampling/neighbor_sampling_sg_v32_e32.cpp
-    src/sampling/neighbor_sampling_sg_v64_e64.cpp
+    src/sampling/neighbor_sampling_mg_v32_e64.cu
+    src/sampling/neighbor_sampling_mg_v32_e32.cu
+    src/sampling/neighbor_sampling_mg_v64_e64.cu
+    src/sampling/neighbor_sampling_sg_v32_e64.cu
+    src/sampling/neighbor_sampling_sg_v32_e32.cu
+    src/sampling/neighbor_sampling_sg_v64_e64.cu
     src/sampling/renumber_sampled_edgelist_sg_v64_e64.cu
     src/sampling/renumber_sampled_edgelist_sg_v32_e32.cu
     src/sampling/sampling_post_processing_sg_v64_e64.cu
diff --git a/cpp/src/community/k_truss_impl.cuh b/cpp/src/community/k_truss_impl.cuh
index e052a892917..4facf093535 100644
--- a/cpp/src/community/k_truss_impl.cuh
+++ b/cpp/src/community/k_truss_impl.cuh
@@ -148,13 +148,13 @@ k_truss(raft::handle_t const& handle,
         std::nullopt,
         std::nullopt,
         cugraph::graph_properties_t{true, graph_view.is_multigraph()},
-        true);
+        false);
 
     modified_graph_view = (*modified_graph).view();
   }
 
   // 2. Find (k-1)-core and exclude edges that do not belong to (k-1)-core
-
+  /*
   {
     auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
 
@@ -211,6 +211,7 @@ k_truss(raft::handle_t const& handle,
 
     renumber_map = std::move(tmp_renumber_map);
   }
+  */
 
   // 3. Keep only the edges from a low-degree vertex to a high-degree vertex.
 
@@ -277,7 +278,7 @@ k_truss(raft::handle_t const& handle,
         std::nullopt,
         std::nullopt,
         cugraph::graph_properties_t{false /* now asymmetric */, cur_graph_view.is_multigraph()},
-        true);
+        false);
 
     modified_graph_view = (*modified_graph).view();
     if (renumber_map) {  // collapse renumber_map
@@ -340,8 +341,12 @@ k_truss(raft::handle_t const& handle,
         edge_weight_view ? std::make_optional(*edge_weight_view) : std::nullopt,
         std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        /*
         std::make_optional(
-          raft::device_span<vertex_t const>((*renumber_map).data(), (*renumber_map).size())));
+          raft::device_span<vertex_t const>((*renumber_map).data(), (*renumber_map).size()))
+        */
+        std::optional<raft::device_span<vertex_t const>>(std::nullopt)
+        );
 
     std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts) =
       symmetrize_edgelist<vertex_t, weight_t, false, multi_gpu>(handle,
@@ -349,6 +354,12 @@ k_truss(raft::handle_t const& handle,
                                                                 std::move(edgelist_dsts),
                                                                 std::move(edgelist_wgts),
                                                                 false);
+    
+    raft::print_device_vector("edgelist_srcs", edgelist_srcs.data(), edgelist_srcs.size(), std::cout);
+    raft::print_device_vector("edgelist_dsts", edgelist_dsts.data(), edgelist_dsts.size(), std::cout);
+
+    printf("\nK-TRUSS Successfully completed and the subgraph size = %d\n", edgelist_srcs.size());
+    //raft::print_device_vector("edgelist_wgts", count_3_.data(), count_3_.size(), std::cout); 
 
     return std::make_tuple(
       std::move(edgelist_srcs), std::move(edgelist_dsts), std::move(edgelist_wgts));
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 3bf569956d8..feaa8824d73 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -16,10 +16,14 @@
 
 #pragma once
 
+#include "prims/fill_edge_property.cuh"
+#include "prims/transform_e.cuh"
+
 #include "sampling/detail/sampling_utils.hpp"
 
 #include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
+
 #include <cugraph/graph.hpp>
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/sampling_functions.hpp>
@@ -70,6 +74,10 @@ neighbor_sample_impl(
   bool do_expensive_check)
 {
 
+  std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
+  std::optional<graph_view_t<vertex_t, edge_t, false, multi_gpu>> modified_graph_view{std::nullopt};
+
+
   if (fan_out) {
     CUGRAPH_EXPECTS((*fan_out).size() > 0,
                     "Invalid input argument: number of levels must be non-zero.");
@@ -92,6 +100,32 @@ neighbor_sample_impl(
           static_cast<size_t>(std::numeric_limits<int32_t>::max()),
       "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                                 // to store hops
+
+    edge_masks_vector.reserve(std::get<0>(*heterogeneous_fan_out).size() - 1);
+    
+    for (int i = 0; i < std::get<0>(*heterogeneous_fan_out).size() - 1; i++) {
+
+      cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
+      
+      cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
+
+      (*modified_graph_view).attach_edge_mask(edge_mask.view());
+
+      cugraph::transform_e(
+        handle,
+        *modified_graph_view,
+        cugraph::edge_src_dummy_property_t{}.view(),
+        cugraph::edge_dst_dummy_property_t{}.view(),
+        *edge_type_view,
+        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
+          return edge_type == valid_edge_type;
+        },
+        edge_mask.mutable_view(),
+        false);
+      
+      edge_masks_vector.push_back(std::move(edge_mask));
+    }
+    
   }
 
   if constexpr (!multi_gpu) {
@@ -165,52 +199,84 @@ neighbor_sample_impl(
 
   std::vector<size_t> level_sizes{};
   int32_t hop{0};
-  for (auto&& k_level : (*fan_out)) {
-    rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
-    rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());
-    std::optional<rmm::device_uvector<weight_t>> weights{std::nullopt};
-    std::optional<rmm::device_uvector<edge_t>> edge_ids{std::nullopt};
-    std::optional<rmm::device_uvector<edge_type_t>> edge_types{std::nullopt};
-    std::optional<rmm::device_uvector<int32_t>> labels{std::nullopt};
-
-    if (k_level > 0) {
-      std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
-        sample_edges(handle,
-                     graph_view,
-                     edge_weight_view,
-                     edge_id_view,
-                     edge_type_view,
-                     edge_bias_view,
-                     rng_state,
-                     this_frontier_vertices,
-                     this_frontier_vertex_labels,
-                     static_cast<size_t>(k_level),
-                     with_replacement);
-    } else {
-      std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
-        gather_one_hop_edgelist(handle,
-                                graph_view,
-                                edge_weight_view,
-                                edge_id_view,
-                                edge_type_view,
-                                this_frontier_vertices,
-                                this_frontier_vertex_labels);
+  int32_t edge_type_id_max{1}; // A value of 1 translate to homogeneous neighbor sample
+  
+  auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
+  while(true) {
+    int32_t k_level{0};
+    if (fan_out) {
+      k_level = (*fan_out)[hop];
+      if (hop == (*fan_out).size()) {
+        break;
+      }
+    } else if (heterogeneous_fan_out) {
+      // initially edge type
+      edge_type_id_max = std::get<0>(*heterogeneous_fan_out).back() - 1;
     }
+  
+    for (int i = 0; i < edge_type_id_max; i++) {
+      if (heterogeneous_fan_out) { 
+        // Can make the code easier to read by setting a mask for both
+        // homogeneous and heterogeneous neighbor sample. For the former,
+        // no edges should be masked
+        cur_graph_view.attach_edge_mask(edge_masks_vector[i].view());
+        auto k_level_size = (std::get<1>(*heterogeneous_fan_out)[i + 1] - std::get<1>(*heterogeneous_fan_out)[i]);
+        if (k_level_size > hop) {
+          k_level = i + hop;
+        } // otherwise, k_level = 0
+      }
+      rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
+      rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> weights{std::nullopt};
+      std::optional<rmm::device_uvector<edge_t>> edge_ids{std::nullopt};
+      std::optional<rmm::device_uvector<edge_type_t>> edge_types{std::nullopt};
+      std::optional<rmm::device_uvector<int32_t>> labels{std::nullopt};
+
+      if (k_level > 0) {
+        std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
+          sample_edges(handle,
+                      cur_graph_view,
+                      edge_weight_view,
+                      edge_id_view,
+                      edge_type_view,
+                      edge_bias_view,
+                      rng_state,
+                      this_frontier_vertices,
+                      this_frontier_vertex_labels,
+                      static_cast<size_t>(k_level),
+                      with_replacement);
+      } else {
+        std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
+          gather_one_hop_edgelist(handle,
+                                  cur_graph_view,
+                                  edge_weight_view,
+                                  edge_id_view,
+                                  edge_type_view,
+                                  this_frontier_vertices,
+                                  this_frontier_vertex_labels);
+      }
 
-    level_sizes.push_back(srcs.size());
+      level_sizes.push_back(srcs.size());
 
-    level_result_src_vectors.push_back(std::move(srcs));
-    level_result_dst_vectors.push_back(std::move(dsts));
-    if (weights) { (*level_result_weight_vectors).push_back(std::move(*weights)); }
-    if (edge_ids) { (*level_result_edge_id_vectors).push_back(std::move(*edge_ids)); }
-    if (edge_types) { (*level_result_edge_type_vectors).push_back(std::move(*edge_types)); }
-    if (labels) { (*level_result_label_vectors).push_back(std::move(*labels)); }
+      level_result_src_vectors.push_back(std::move(srcs));
+      level_result_dst_vectors.push_back(std::move(dsts));
+      if (weights) { (*level_result_weight_vectors).push_back(std::move(*weights)); }
+      if (edge_ids) { (*level_result_edge_id_vectors).push_back(std::move(*edge_ids)); }
+      if (edge_types) { (*level_result_edge_type_vectors).push_back(std::move(*edge_types)); }
+      if (labels) { (*level_result_label_vectors).push_back(std::move(*labels)); }
+    }
 
     ++hop;
+
+
+    // aggregation here for each hope regardless of the type  ****************************************************
+    // Genrate the mask for each type and store it for each iteration. mask vector at the beginning, populate that
+    // and within the inner loop, activate the right mask based and proceed.
+    
     if (hop < (*fan_out).size()) {
       // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span
       //  rather than making a copy.
-      auto vertex_partition_range_lasts = graph_view.vertex_partition_range_lasts();
+      auto vertex_partition_range_lasts = cur_graph_view.vertex_partition_range_lasts();
       std::tie(frontier_vertices, frontier_vertex_labels, vertex_used_as_source) =
         prepare_next_frontier(
           handle,
@@ -223,7 +289,7 @@ neighbor_sample_impl(
                                      level_result_label_vectors->back().size()))
                                  : std::nullopt,
           std::move(vertex_used_as_source),
-          graph_view.local_vertex_partition_view(),
+          cur_graph_view.local_vertex_partition_view(),
           vertex_partition_range_lasts,
           prior_sources_behavior,
           dedupe_sources,
@@ -237,7 +303,7 @@ neighbor_sample_impl(
           frontier_vertex_labels->data(), frontier_vertex_labels->size());
       }
     }
-  }
+  } // While loop end here: FIXME: remove this comment
 
   auto result_size = std::reduce(level_sizes.begin(), level_sizes.end());
   size_t output_offset{};
@@ -351,7 +417,6 @@ neighbor_sample_impl(
                                              std::move(result_hops),
                                              std::move(result_labels),
                                              label_to_output_comm_rank);
-#endif
 }
 
 }  // namespace detail
@@ -501,6 +566,235 @@ heterogeneous_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check)
 {
+  // for loop : over the number of edgetypes /
+  // inside the loop, create an edge mask and update the edge mask to mask out all the irrelevant edges
+  // adn call neighborhood sampling
+
+  // before this, there is a reordering constrain.
+  // 10 labels amd 100 seeds, 3 edge type. For each seed, you provide sample of each edge type (0, 1, 2)
+  // soritng criterial, label , seed, edge tyope
+  // label, edge type, seed
+
+  // Pick how to sort - requries assembling? if the natural output order based on this scheme if it doesn't meet the requirement, needs to shuffle
+  // shuffling, label to be the primiary key? if you loop over the edge tyoe, edge type will be the primary key.
+  
+  // Goal: reorder the sample at the end - Label is the primary key? then label? then seed?
+  // stable sort? to not mess up the internal order? "Use scatter operation to be more efficient?"
+
+  // Iterate over all edge types
+
+  cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
+  cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
+
+
+  std::vector<rmm::device_uvector<vertex_t>> level_result_src_vectors{};
+  std::vector<rmm::device_uvector<vertex_t>> level_result_dst_vectors{};
+  auto level_result_weight_vectors =
+    edge_weight_view ? std::make_optional(std::vector<rmm::device_uvector<weight_t>>{})
+                     : std::nullopt;
+  auto level_result_edge_id_vectors =
+    edge_id_view ? std::make_optional(std::vector<rmm::device_uvector<edge_t>>{}) : std::nullopt;
+  auto level_result_edge_type_vectors =
+    edge_type_view ? std::make_optional(std::vector<rmm::device_uvector<edge_type_t>>{})
+                   : std::nullopt;
+  auto level_result_label_vectors =
+    starting_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
+                                : std::nullopt;
+  
+    level_result_src_vectors.reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
+  level_result_dst_vectors.reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
+  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
+  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
+  if (level_result_edge_type_vectors) {
+    (*level_result_edge_type_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
+  }
+  if (level_result_label_vectors) { (*level_result_label_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
+
+
+  
+
+
+
+
+
+
+
+
+  std::vector<size_t> level_sizes{};
+  for (int i = 0; i < std::get<0>(heterogeneous_fan_out).size() - 1; i++) {
+
+    cugraph::transform_e(
+        handle,
+        graph_view,
+        cugraph::edge_src_dummy_property_t{}.view(),
+        cugraph::edge_dst_dummy_property_t{}.view(),
+        *edge_type_view,
+        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
+          return edge_type == valid_edge_type;
+        },
+        edge_mask.mutable_view(),
+        false);
+    
+    auto [
+      level_result_srcs,
+      level_result_dsts,
+      level_result_weights,
+      level_result_edge_ids,
+      level_result_edge_types,
+      level_result_hops,
+      level_result_labels,
+      label_to_output_comm_rank_ // unused
+      ] = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+        handle,
+        graph_view,
+        edge_weight_view,
+        edge_id_view,
+        edge_type_view,
+        edge_bias_view,
+        starting_vertices,
+        starting_vertex_labels,
+        label_to_output_comm_rank,
+        std::nullopt,
+        heterogeneous_fan_out,
+        return_hops,
+        with_replacement,
+        prior_sources_behavior,
+        dedupe_sources,
+        rng_state,
+        do_expensive_check);
+    
+    level_sizes.push_back(level_result_srcs.size());
+
+    level_result_src_vectors.push_back(std::move(level_result_srcs));
+    level_result_dst_vectors.push_back(std::move(level_result_dsts));
+    if (level_result_weights) { (*level_result_weight_vectors).push_back(std::move(*level_result_weights)); }
+    if (level_result_edge_ids) { (*level_result_edge_id_vectors).push_back(std::move(*level_result_edge_ids)); }
+    if (level_result_edge_types) { (*level_result_edge_type_vectors).push_back(std::move(*level_result_edge_types)); }
+    if (level_result_labels) { (*level_result_label_vectors).push_back(std::move(*level_result_labels)); }
+  }
+
+
+  auto result_size = std::reduce(level_sizes.begin(), level_sizes.end());
+
+  size_t output_offset{};
+
+  rmm::device_uvector<vertex_t> result_srcs(result_size, handle.get_stream());
+  output_offset = 0;
+  for (size_t i = 0; i < level_result_src_vectors.size(); ++i) {
+    raft::copy(result_srcs.begin() + output_offset,
+               level_result_src_vectors[i].begin(),
+               level_sizes[i],
+               handle.get_stream());
+    output_offset += level_sizes[i];
+  }
+  level_result_src_vectors.clear();
+  level_result_src_vectors.shrink_to_fit();
+
+  rmm::device_uvector<vertex_t> result_dsts(result_size, handle.get_stream());
+  output_offset = 0;
+  for (size_t i = 0; i < level_result_dst_vectors.size(); ++i) {
+    raft::copy(result_dsts.begin() + output_offset,
+               level_result_dst_vectors[i].begin(),
+               level_sizes[i],
+               handle.get_stream());
+    output_offset += level_sizes[i];
+  }
+  level_result_dst_vectors.clear();
+  level_result_dst_vectors.shrink_to_fit();
+
+  auto result_weights =
+    level_result_weight_vectors
+      ? std::make_optional(rmm::device_uvector<weight_t>(result_size, handle.get_stream()))
+      : std::nullopt;
+  if (result_weights) {
+    output_offset = 0;
+    for (size_t i = 0; i < (*level_result_weight_vectors).size(); ++i) {
+      raft::copy((*result_weights).begin() + output_offset,
+                 (*level_result_weight_vectors)[i].begin(),
+                 level_sizes[i],
+                 handle.get_stream());
+      output_offset += level_sizes[i];
+    }
+    level_result_weight_vectors = std::nullopt;
+  }
+
+  auto result_edge_ids =
+    level_result_edge_id_vectors
+      ? std::make_optional(rmm::device_uvector<edge_t>(result_size, handle.get_stream()))
+      : std::nullopt;
+  if (result_edge_ids) {
+    output_offset = 0;
+    for (size_t i = 0; i < (*level_result_edge_id_vectors).size(); ++i) {
+      raft::copy((*result_edge_ids).begin() + output_offset,
+                 (*level_result_edge_id_vectors)[i].begin(),
+                 level_sizes[i],
+                 handle.get_stream());
+      output_offset += level_sizes[i];
+    }
+    level_result_edge_id_vectors = std::nullopt;
+  }
+
+  auto result_edge_types =
+    level_result_edge_type_vectors
+      ? std::make_optional(rmm::device_uvector<edge_type_t>(result_size, handle.get_stream()))
+      : std::nullopt;
+  if (result_edge_types) {
+    output_offset = 0;
+    for (size_t i = 0; i < (*level_result_edge_type_vectors).size(); ++i) {
+      raft::copy((*result_edge_types).begin() + output_offset,
+                 (*level_result_edge_type_vectors)[i].begin(),
+                 level_sizes[i],
+                 handle.get_stream());
+      output_offset += level_sizes[i];
+    }
+    level_result_edge_type_vectors = std::nullopt;
+  }
+
+  std::optional<rmm::device_uvector<int32_t>> result_hops{std::nullopt};
+
+  if (return_hops) {
+    result_hops   = rmm::device_uvector<int32_t>(result_size, handle.get_stream());
+    output_offset = 0;
+    for (size_t i = 0; i < std::get<0>(heterogeneous_fan_out).size() - 1; ++i) {
+      detail::scalar_fill(
+        handle, result_hops->data() + output_offset, level_sizes[i], static_cast<int32_t>(i));
+      output_offset += level_sizes[i];
+    }
+  }
+
+  auto result_labels =
+    level_result_label_vectors
+      ? std::make_optional(rmm::device_uvector<label_t>(result_size, handle.get_stream()))
+      : std::nullopt;
+  if (result_labels) {
+    output_offset = 0;
+    for (size_t i = 0; i < (*level_result_label_vectors).size(); ++i) {
+      raft::copy((*result_labels).begin() + output_offset,
+                 (*level_result_label_vectors)[i].begin(),
+                 level_sizes[i],
+                 handle.get_stream());
+      output_offset += level_sizes[i];
+    }
+    level_result_label_vectors = std::nullopt;
+  }
+
+  // 
+
+  /*
+  return std::make_tuple(
+    std::move(result_srcs),
+    std::move(result_dsts),
+    std::move(result_weights),
+    std::move(result_edge_ids),
+    std::move(result_edge_types),
+    std::move(result_hops),
+    std::move(result_labels),
+    //offsets
+    //std::move(label_to_output_comm_rank)
+  );
+  */
+
+  // Remove and pick the one on top
   return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
     graph_view,
@@ -519,8 +813,37 @@ heterogeneous_neighbor_sample(
     dedupe_sources,
     rng_state,
     do_expensive_check);
+
+
+
+
+
+
+
 }
 
+  /*
+  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    edge_bias_view,
+    starting_vertices,
+    starting_vertex_labels,
+    label_to_output_comm_rank,
+    std::nullopt,
+    heterogeneous_fan_out,
+    return_hops,
+    with_replacement,
+    prior_sources_behavior,
+    dedupe_sources,
+    rng_state,
+    do_expensive_check);
+//}
+*/
+
 
 
 
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_mg_v32_e32.cpp
rename to cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_mg_v32_e64.cpp
rename to cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_mg_v64_e64.cpp
rename to cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_sg_v32_e32.cpp
rename to cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_sg_v32_e64.cpp
rename to cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
similarity index 100%
rename from cpp/src/sampling/neighbor_sampling_sg_v64_e64.cpp
rename to cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
diff --git a/cpp/tests/community/k_truss_test.cpp b/cpp/tests/community/k_truss_test.cpp
index 424d52f2067..af34958f0de 100644
--- a/cpp/tests/community/k_truss_test.cpp
+++ b/cpp/tests/community/k_truss_test.cpp
@@ -262,7 +262,6 @@ class Tests_KTruss : public ::testing::TestWithParam<std::tuple<KTruss_Usecase,
 };
 
 using Tests_KTruss_File = Tests_KTruss<cugraph::test::File_Usecase>;
-using Tests_KTruss_Rmat = Tests_KTruss<cugraph::test::Rmat_Usecase>;
 
 TEST_P(Tests_KTruss_File, CheckInt32Int32Float)
 {
@@ -270,57 +269,16 @@ TEST_P(Tests_KTruss_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_KTruss_File, CheckInt64Int64Float)
-{
-  run_current_test<int64_t, int64_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
-TEST_P(Tests_KTruss_Rmat, CheckInt32Int32Float)
-{
-  run_current_test<int32_t, int32_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
-TEST_P(Tests_KTruss_Rmat, CheckInt64Int64Float)
-{
-  run_current_test<int64_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 INSTANTIATE_TEST_SUITE_P(
   simple_test,
   Tests_KTruss_File,
   ::testing::Combine(
     // enable correctness checks
-    ::testing::Values(KTruss_Usecase{5, true, false},
-                      KTruss_Usecase{4, true, false},
-                      KTruss_Usecase{9, true, true},
-                      KTruss_Usecase{7, true, true}),
-    ::testing::Values(cugraph::test::File_Usecase("test/datasets/netscience.mtx"),
-                      cugraph::test::File_Usecase("test/datasets/dolphins.mtx"))));
-
-INSTANTIATE_TEST_SUITE_P(rmat_small_test,
-                         Tests_KTruss_Rmat,
-                         // enable correctness checks
-                         ::testing::Combine(::testing::Values(KTruss_Usecase{5, false, true},
-                                                              KTruss_Usecase{4, false, true},
-                                                              KTruss_Usecase{9, true, true},
-                                                              KTruss_Usecase{7, true, true}),
-                                            ::testing::Values(cugraph::test::Rmat_Usecase(
-                                              10, 16, 0.57, 0.19, 0.19, 0, true, false))));
-
-INSTANTIATE_TEST_SUITE_P(
-  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
-                          --gtest_filter to select only the rmat_benchmark_test with a specific
-                          vertex & edge type combination) by command line arguments and do not
-                          include more than one Rmat_Usecase that differ only in scale or edge
-                          factor (to avoid running same benchmarks more than once) */
-  Tests_KTruss_Rmat,
-  // disable correctness checks for large graphs
-  // FIXME: High memory footprint. Perform nbr_intersection in chunks.
-  ::testing::Combine(
-    ::testing::Values(KTruss_Usecase{12, false, false}),
-    ::testing::Values(cugraph::test::Rmat_Usecase(14, 16, 0.57, 0.19, 0.19, 0, true, false))));
+    ::testing::Values(KTruss_Usecase{5, false, false}),
+    ::testing::Values(
+      cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/datasets/test_datasets_.mtx")
+      //cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/cugraph/datasets/dolphins.mtx"),
+      //cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/cugraph/datasets/netscience.mtx")
+      )));
 
 CUGRAPH_TEST_PROGRAM_MAIN()

From fe35c805c276e4ec8ef7182f4fa7fd9dc465f718 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 9 Sep 2024 16:42:33 -0700
Subject: [PATCH 031/117] add exit condition

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index feaa8824d73..bde467b9407 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -200,8 +200,14 @@ neighbor_sample_impl(
   std::vector<size_t> level_sizes{};
   int32_t hop{0};
   int32_t edge_type_id_max{1}; // A value of 1 translate to homogeneous neighbor sample
+  int32_t num_edge_type_per_hop{0};
   
   auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
+
+  if (heterogeneous_fan_out) {
+    num_edge_type_per_hop = std::get<0>(*heterogeneous_fan_out).back() - 1;
+  }
+
   while(true) {
     int32_t k_level{0};
     if (fan_out) {
@@ -210,8 +216,11 @@ neighbor_sample_impl(
         break;
       }
     } else if (heterogeneous_fan_out) {
-      // initially edge type
-      edge_type_id_max = std::get<0>(*heterogeneous_fan_out).back() - 1;
+        if (num_edge_type_per_hop == 0) {
+          break;
+        }
+        edge_type_id_max = std::get<0>(*heterogeneous_fan_out).back() - 1;
+      
     }
   
     for (int i = 0; i < edge_type_id_max; i++) {
@@ -223,7 +232,10 @@ neighbor_sample_impl(
         auto k_level_size = (std::get<1>(*heterogeneous_fan_out)[i + 1] - std::get<1>(*heterogeneous_fan_out)[i]);
         if (k_level_size > hop) {
           k_level = i + hop;
-        } // otherwise, k_level = 0
+        } else { // otherwise, k_level = 0
+          --num_edge_type_per_hop ;
+
+        }
       }
       rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
       rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());

From a658b29d0cb1585c958f4cd171a4a8688eeab573 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 10 Sep 2024 10:02:43 -0700
Subject: [PATCH 032/117] remove comments

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index bde467b9407..3c24aaa1dc7 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -279,11 +279,6 @@ neighbor_sample_impl(
     }
 
     ++hop;
-
-
-    // aggregation here for each hope regardless of the type  ****************************************************
-    // Genrate the mask for each type and store it for each iteration. mask vector at the beginning, populate that
-    // and within the inner loop, activate the right mask based and proceed.
     
     if (hop < (*fan_out).size()) {
       // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span

From e52a38a42adecc1777abd5044adcc25195a4bd57 Mon Sep 17 00:00:00 2001
From: Charles Hastings <charlesh@nvidia.com>
Date: Wed, 11 Sep 2024 16:32:33 -0700
Subject: [PATCH 033/117] Add Implementation

---
 cpp/CMakeLists.txt                            |   1 +
 .../sampling/detail/conversion_utilities.cu   |  62 +++++
 .../detail/conversion_utilities_impl.cuh      |  58 +++++
 cpp/src/sampling/detail/sampling_utils.hpp    |  36 ++-
 .../shuffle_and_organize_output_impl.cuh      |  57 ++---
 .../shuffle_and_organize_output_mg_v32_e32.cu |   6 +-
 .../shuffle_and_organize_output_mg_v32_e64.cu |   6 +-
 .../shuffle_and_organize_output_mg_v64_e64.cu |   6 +-
 cpp/src/sampling/neighbor_sampling_impl.hpp   | 211 ++++++++++++++----
 9 files changed, 343 insertions(+), 100 deletions(-)
 create mode 100644 cpp/src/sampling/detail/conversion_utilities.cu
 create mode 100644 cpp/src/sampling/detail/conversion_utilities_impl.cuh

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index b8eaba9d575..7849bedce44 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -188,6 +188,7 @@ set(CUGRAPH_SOURCES
     src/detail/groupby_and_count_mg_v64_e64.cu
     src/detail/collect_comm_wrapper_mg_v32_e32.cu
     src/detail/collect_comm_wrapper_mg_v64_e64.cu
+    src/sampling/detail/conversion_utilities.cu
     src/sampling/random_walks_mg_v64_e64.cu
     src/sampling/random_walks_mg_v32_e32.cu
     src/sampling/random_walks_mg_v32_e64.cu
diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
new file mode 100644
index 00000000000..5e6f4b00fd2
--- /dev/null
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "sampling/detail/sampling_utils.hpp"
+
+#include <cugraph/utilities/misc_utils.cuh>
+
+namespace cugraph {
+namespace detail {
+
+rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets)
+{
+  return expand_sparse_offsets(starting_vertex_offsets, int32_t{0}, handle.get_stream());
+}
+
+template <typename label_t>
+rmm::device_uvector<int32_t> flatten_label_map(
+  raft::handle_t const& handle,
+  std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>
+    label_to_output_comm_rank)
+{
+  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
+
+  label_t max_label = thrust::reduce(handle.get_thrust_policy(),
+                                     std::get<0>(label_to_output_comm_rank).begin(),
+                                     std::get<0>(label_to_output_comm_rank).end(),
+                                     label_t{0},
+                                     thrust::maximum<label_t>());
+
+  label_map.resize(max_label, handle.get_stream());
+
+  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), label_t{0});
+  thrust::gather(handle.get_thrust_policy(),
+                 std::get<0>(label_to_output_comm_rank).begin(),
+                 std::get<0>(label_to_output_comm_rank).end(),
+                 std::get<1>(label_to_output_comm_rank).begin(),
+                 label_map.begin());
+
+  return label_map;
+}
+
+template rmm::device_uvector<int32_t> flatten_label_map(
+  raft::handle_t const& handle,
+  std::tuple<raft::device_span<int32_t const>, raft::device_span<int32_t const>>
+    label_to_output_comm_rank);
+
+}  // namespace detail
+}  // namespace cugraph
diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
new file mode 100644
index 00000000000..3ed11d05181
--- /dev/null
+++ b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "sampling/detail/sampling_utils.hpp"
+
+#include <cugraph/utilities/misc_utils.cuh>
+
+namespace cugraph {
+namespace detail {
+
+rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets)
+{
+  return expand_sparse_offsets(*starting_vertex_offsets, int32_t{0}, handle.get_stream());
+}
+
+rmm::device_uvector<int32_t> flatten_label_map(
+  raft::handle_t const& handle,
+  std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>
+    label_to_output_comm_rank)
+{
+  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
+
+  label_t max_label = thrust::reduce(handle.get_thrust_policy(),
+                                     std::get<0>(label_to_output_comm_rank).begin(),
+                                     std::get<0>(label_to_output_comm_rank).end(),
+                                     label_t{0},
+                                     thrust::maximum<label_t>());
+
+  label_map.resize(max_label, handle.get_stream());
+
+  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), label_t{0});
+  thrust::gather(handle.get_thrust_policy(),
+                 std::get<0>(label_to_output_comm_rank).begin(),
+                 std::get<0>(label_to_output_comm_rank).end(),
+                 std::get<1>(label_to_output_comm_rank).begin(),
+                 label_map.begin());
+
+  return label_map;
+}
+
+}  // namespace detail
+}  // namespace cugraph
diff --git a/cpp/src/sampling/detail/sampling_utils.hpp b/cpp/src/sampling/detail/sampling_utils.hpp
index 102f9ec58f7..b68f3fc74f7 100644
--- a/cpp/src/sampling/detail/sampling_utils.hpp
+++ b/cpp/src/sampling/detail/sampling_utils.hpp
@@ -293,7 +293,41 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<edge_type_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<label_t>>&& labels,
-  std::optional<std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>>
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank);
+
+/**
+ * @brief   Convert the starting vertex offsets into starting vertex labels
+ *
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param starting_vertex_offsets Offsets array defining where each vertex label begins
+ *
+ * @returns device vector containing labels for each starting vertex
+ */
+rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets);
+
+/**
+ * @brief   Flatten the legacy label_to_output_comm_rank into the new structure
+ *
+ * Legacy structure supported arbitrary labels, the new structure is a dense mapping of labels from
+ * [0,n).
+ *
+ * @tparam label_t typename for the label
+ *
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param label_to_output_comm_rank  A tuple containing label ids and the comm rank each label
+ * should be assigned to
+ *
+ * @returns device vector containing the mapping to comm_rank.  Entry `i` will be the comm rank
+ * destination for label `i`.
+ */
+template <typename label_t>
+rmm::device_uvector<int32_t> flatten_label_map(
+  raft::handle_t const& handle,
+  std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>
     label_to_output_comm_rank);
+
 }  // namespace detail
 }  // namespace cugraph
diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_impl.cuh b/cpp/src/sampling/detail/shuffle_and_organize_output_impl.cuh
index ec14e99baec..391dd99b1df 100644
--- a/cpp/src/sampling/detail/shuffle_and_organize_output_impl.cuh
+++ b/cpp/src/sampling/detail/shuffle_and_organize_output_impl.cuh
@@ -41,14 +41,12 @@ namespace detail {
 
 template <typename label_t>
 struct shuffle_to_output_comm_rank_t {
-  raft::device_span<label_t const> output_label_;
   raft::device_span<int32_t const> output_rank_;
 
   template <typename key_t>
   __device__ int32_t operator()(key_t key) const
   {
-    auto pos = thrust::lower_bound(thrust::seq, output_label_.begin(), output_label_.end(), key);
-    return output_rank_[thrust::distance(output_label_.begin(), pos)];
+    return output_rank_[key];
   }
 };
 
@@ -206,8 +204,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<edge_type_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<label_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank)
 {
   std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
 
@@ -215,8 +212,6 @@ shuffle_and_organize_output(
     sort_sampled_tuples(handle, majors, minors, weights, edge_ids, edge_types, hops, *labels);
 
     if (label_to_output_comm_rank) {
-      CUGRAPH_EXPECTS(labels, "labels must be specified in order to shuffle sampling results");
-
       auto& comm           = handle.get_comms();
       auto const comm_size = comm.get_size();
 
@@ -247,8 +242,7 @@ shuffle_and_organize_output(
                                           edge_ids->begin(),
                                           edge_types->begin(),
                                           hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -282,8 +276,7 @@ shuffle_and_organize_output(
                                           weights->begin(),
                                           edge_ids->begin(),
                                           edge_types->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -317,8 +310,7 @@ shuffle_and_organize_output(
                                           weights->begin(),
                                           edge_ids->begin(),
                                           hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -347,8 +339,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), weights->begin(), edge_ids->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -383,8 +374,7 @@ shuffle_and_organize_output(
                                           weights->begin(),
                                           edge_types->begin(),
                                           hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -413,8 +403,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), weights->begin(), edge_types->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -444,8 +433,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), weights->begin(), hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -471,8 +459,7 @@ shuffle_and_organize_output(
                 labels->begin(),
                 labels->end(),
                 thrust::make_zip_iterator(majors.begin(), minors.begin(), weights->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -505,8 +492,7 @@ shuffle_and_organize_output(
                                           edge_ids->begin(),
                                           edge_types->begin(),
                                           hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -535,8 +521,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), edge_ids->begin(), edge_types->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -566,8 +551,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), edge_ids->begin(), hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -593,8 +577,7 @@ shuffle_and_organize_output(
                 labels->begin(),
                 labels->end(),
                 thrust::make_zip_iterator(majors.begin(), minors.begin(), edge_ids->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -623,8 +606,7 @@ shuffle_and_organize_output(
                 labels->end(),
                 thrust::make_zip_iterator(
                   majors.begin(), minors.begin(), edge_types->begin(), hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -651,8 +633,7 @@ shuffle_and_organize_output(
                 labels->begin(),
                 labels->end(),
                 thrust::make_zip_iterator(majors.begin(), minors.begin(), edge_types->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -678,8 +659,7 @@ shuffle_and_organize_output(
                 labels->begin(),
                 labels->end(),
                 thrust::make_zip_iterator(majors.begin(), minors.begin(), hops->begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
@@ -702,8 +682,7 @@ shuffle_and_organize_output(
                 labels->begin(),
                 labels->end(),
                 thrust::make_zip_iterator(majors.begin(), minors.begin()),
-                shuffle_to_output_comm_rank_t<label_t>{std::get<0>(*label_to_output_comm_rank),
-                                                       std::get<1>(*label_to_output_comm_rank)},
+                shuffle_to_output_comm_rank_t<label_t>{*label_to_output_comm_rank},
                 comm_size,
                 mem_frugal_threshold,
                 handle.get_stream());
diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu
index 73a152487ca..4a264469c97 100644
--- a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu
+++ b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu
@@ -36,8 +36,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
@@ -56,8 +55,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 }  // namespace detail
 }  // namespace cugraph
diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
index 50fc3910aca..ef760844b51 100644
--- a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
+++ b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
@@ -36,8 +36,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
@@ -56,8 +55,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 }  // namespace detail
 }  // namespace cugraph
diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu
index ff7a716e609..f66ce3e2d63 100644
--- a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu
+++ b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu
@@ -36,8 +36,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
@@ -56,8 +55,7 @@ shuffle_and_organize_output(
   std::optional<rmm::device_uvector<int32_t>>&& edge_types,
   std::optional<rmm::device_uvector<int32_t>>&& hops,
   std::optional<rmm::device_uvector<int32_t>>&& 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::device_span<int32_t const>> label_to_output_comm_rank);
 
 }  // namespace detail
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 1785f934426..ef8fbc46325 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -48,31 +48,23 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<label_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-neighbor_sample_impl(
-  raft::handle_t const& handle,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
-  raft::device_span<vertex_t const> this_frontier_vertices,
-  std::optional<raft::device_span<label_t const>> this_frontier_vertex_labels,
-  std::optional<std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>>
-    label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  raft::random::RngState& rng_state,
-  bool do_expensive_check)
+neighbor_sample_impl(raft::handle_t const& handle,
+                     raft::random::RngState& rng_state,
+                     graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+                     std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+                     std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+                     std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+                     std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+                     raft::device_span<vertex_t const> starting_vertices,
+                     std::optional<raft::device_span<label_t const>> starting_vertex_labels,
+                     std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+                     raft::host_span<int32_t const> fan_out,
+                     bool return_hops,
+                     bool with_replacement,
+                     prior_sources_behavior_t prior_sources_behavior,
+                     bool dedupe_sources,
+                     bool do_expensive_check)
 {
-#ifdef NO_CUGRAPH_OPS  // FIXME: this is relevant only when edge_bias_view.has_value() is false,
-                       // this ifdef statement will be removed once we migrate relevant cugraph-ops
-                       // functions to cugraph
-  CUGRAPH_FAIL(
-    "neighbor_sample_impl not supported in this configuration, built with NO_CUGRAPH_OPS");
-#else
   static_assert(std::is_floating_point_v<bias_t>);
 
   CUGRAPH_EXPECTS(fan_out.size() > 0, "Invalid input argument: number of levels must be non-zero.");
@@ -87,8 +79,8 @@ neighbor_sample_impl(
   }
 
   CUGRAPH_EXPECTS(
-    !label_to_output_comm_rank || this_frontier_vertex_labels,
-    "cannot specify output GPU mapping without also specifying this_frontier_vertex_labels");
+    !label_to_output_comm_rank || starting_vertex_labels,
+    "cannot specify output GPU mapping without also specifying starting_vertex_labels");
 
   if (do_expensive_check) {
     if (edge_bias_view) {
@@ -102,11 +94,6 @@ neighbor_sample_impl(
                       "Invalid input argument: sum of neighboring edge bias values should not "
                       "exceed std::numeric_limits<bias_t>::max() for any vertex.");
     }
-
-    if (label_to_output_comm_rank) {
-      CUGRAPH_EXPECTS(cugraph::detail::is_sorted(handle, std::get<0>(*label_to_output_comm_rank)),
-                      "Labels in label_to_output_comm_rank must be sorted");
-    }
   }
 
   std::vector<rmm::device_uvector<vertex_t>> level_result_src_vectors{};
@@ -120,8 +107,8 @@ neighbor_sample_impl(
     edge_type_view ? std::make_optional(std::vector<rmm::device_uvector<edge_type_t>>{})
                    : std::nullopt;
   auto level_result_label_vectors =
-    this_frontier_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
-                                : std::nullopt;
+    starting_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
+                           : std::nullopt;
 
   level_result_src_vectors.reserve(fan_out.size());
   level_result_dst_vectors.reserve(fan_out.size());
@@ -132,7 +119,7 @@ neighbor_sample_impl(
 
   rmm::device_uvector<vertex_t> frontier_vertices(0, handle.get_stream());
   auto frontier_vertex_labels =
-    this_frontier_vertex_labels
+    starting_vertex_labels
       ? std::make_optional(rmm::device_uvector<label_t>{0, handle.get_stream()})
       : std::nullopt;
 
@@ -143,7 +130,7 @@ neighbor_sample_impl(
   if (prior_sources_behavior == prior_sources_behavior_t::EXCLUDE) {
     vertex_used_as_source = std::make_optional(
       std::make_tuple(rmm::device_uvector<vertex_t>{0, handle.get_stream()},
-                      this_frontier_vertex_labels
+                      starting_vertex_labels
                         ? std::make_optional(rmm::device_uvector<label_t>{0, handle.get_stream()})
                         : std::nullopt));
   }
@@ -167,8 +154,8 @@ neighbor_sample_impl(
                      edge_type_view,
                      edge_bias_view,
                      rng_state,
-                     this_frontier_vertices,
-                     this_frontier_vertex_labels,
+                     starting_vertices,
+                     starting_vertex_labels,
                      static_cast<size_t>(k_level),
                      with_replacement);
     } else {
@@ -178,8 +165,8 @@ neighbor_sample_impl(
                                 edge_weight_view,
                                 edge_id_view,
                                 edge_type_view,
-                                this_frontier_vertices,
-                                this_frontier_vertex_labels);
+                                starting_vertices,
+                                starting_vertex_labels);
     }
 
     level_sizes.push_back(srcs.size());
@@ -199,8 +186,8 @@ neighbor_sample_impl(
       std::tie(frontier_vertices, frontier_vertex_labels, vertex_used_as_source) =
         prepare_next_frontier(
           handle,
-          this_frontier_vertices,
-          this_frontier_vertex_labels,
+          starting_vertices,
+          starting_vertex_labels,
           raft::device_span<vertex_t const>{level_result_dst_vectors.back().data(),
                                             level_result_dst_vectors.back().size()},
           frontier_vertex_labels ? std::make_optional(raft::device_span<label_t const>(
@@ -214,12 +201,12 @@ neighbor_sample_impl(
           dedupe_sources,
           do_expensive_check);
 
-      this_frontier_vertices =
+      starting_vertices =
         raft::device_span<vertex_t const>(frontier_vertices.data(), frontier_vertices.size());
 
       if (frontier_vertex_labels) {
-        this_frontier_vertex_labels = raft::device_span<label_t const>(
-          frontier_vertex_labels->data(), frontier_vertex_labels->size());
+        starting_vertex_labels = raft::device_span<label_t const>(frontier_vertex_labels->data(),
+                                                                  frontier_vertex_labels->size());
       }
     }
   }
@@ -336,7 +323,6 @@ neighbor_sample_impl(
                                              std::move(result_hops),
                                              std::move(result_labels),
                                              label_to_output_comm_rank);
-#endif
 }
 
 }  // namespace detail
@@ -375,8 +361,16 @@ uniform_neighbor_sample(
   bool do_expensive_check)
 {
   using bias_t = weight_t;  // dummy
+
+  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
+
+  if (label_to_output_comm_rank) {
+    label_map = detail::flatten_label_map(handle, *label_to_output_comm_rank);
+  }
+
   return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
+    rng_state,
     graph_view,
     edge_weight_view,
     edge_id_view,
@@ -384,13 +378,14 @@ uniform_neighbor_sample(
     std::nullopt,
     starting_vertices,
     starting_vertex_labels,
-    label_to_output_comm_rank,
+    label_to_output_comm_rank
+      ? std::make_optional(raft::device_span<int32_t const>{label_map.data(), label_map.size()})
+      : std::nullopt,
     fan_out,
     return_hops,
     with_replacement,
     prior_sources_behavior,
     dedupe_sources,
-    rng_state,
     do_expensive_check);
 }
 
@@ -429,8 +424,15 @@ biased_neighbor_sample(
   bool dedupe_sources,
   bool do_expensive_check)
 {
+  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
+
+  if (label_to_output_comm_rank) {
+    label_map = detail::flatten_label_map(handle, *label_to_output_comm_rank);
+  }
+
   return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
+    rng_state,
     graph_view,
     edge_weight_view,
     edge_id_view,
@@ -438,13 +440,126 @@ biased_neighbor_sample(
     edge_bias_view,
     starting_vertices,
     starting_vertex_labels,
-    label_to_output_comm_rank,
+    label_to_output_comm_rank
+      ? std::make_optional(raft::device_span<int32_t const>{label_map.data(), label_map.size()})
+      : std::nullopt,
     fan_out,
     return_hops,
     with_replacement,
     prior_sources_behavior,
     dedupe_sources,
+    do_expensive_check);
+}
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check)
+{
+  using bias_t = weight_t;  // dummy
+
+  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
+
+  if (starting_vertex_offsets)
+    starting_vertex_labels =
+      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
+
+  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
+    rng_state,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    std::nullopt,
+    starting_vertices,
+    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
+                                starting_vertex_labels.data(), starting_vertex_labels.size()})
+                            : std::nullopt,
+    label_to_output_comm_rank,
+    fan_out,
+    sampling_flags.return_hops,
+    sampling_flags.with_replacement,
+    sampling_flags.prior_sources_behavior,
+    sampling_flags.dedupe_sources,
+    do_expensive_check);
+}
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          typename bias_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check)
+{
+  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
+
+  if (starting_vertex_offsets)
+    starting_vertex_labels =
+      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
+
+  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
     rng_state,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    edge_bias_view,
+    starting_vertices,
+    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
+                                starting_vertex_labels.data(), starting_vertex_labels.size()})
+                            : std::nullopt,
+    label_to_output_comm_rank,
+    fan_out,
+    sampling_flags.return_hops,
+    sampling_flags.with_replacement,
+    sampling_flags.prior_sources_behavior,
+    sampling_flags.dedupe_sources,
     do_expensive_check);
 }
 

From c4164397d8cde120c1de01118fd3828111042eb1 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 12 Sep 2024 18:45:53 -0700
Subject: [PATCH 034/117] call heterogeneous renumbering

---
 cpp/include/cugraph/sampling_functions.hpp |   2 +-
 cpp/src/c_api/neighbor_sampling.cpp        | 212 +++++++++++++--------
 2 files changed, 130 insertions(+), 84 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 45114e06480..8f14abca7b4 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -247,7 +247,7 @@ biased_neighbor_sample(
   bool do_expensive_check                         = false);
 
 /**
- * @brief Homogeneous Neighborhood Sampling.
+ * @brief Heterogeneous Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index def04c021e7..218b3e17214 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -325,93 +325,139 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                           (options_.compression_type_ == cugraph_compression_type_t::COO);
 
       if (options_.renumber_results_) {
-        if (options_.compression_type_ == cugraph_compression_type_t::COO) {
-          // COO
+        if (fan_out_ != nullptr) {
+          if (options_.compression_type_ == cugraph_compression_type_t::COO) {
+            // COO
+
+            rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
+            rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+            std::tie(output_majors,
+                    minors,
+                    wgt,
+                    edge_id,
+                    edge_type,
+                    label_hop_offsets,
+                    output_renumber_map,
+                    renumber_map_offsets) =
+              cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
+                handle_,
+                std::move(src),
+                std::move(dst),
+                std::move(wgt),
+                std::move(edge_id),
+                std::move(edge_type),
+                std::move(hop),
+                options_.retain_seeds_
+                  ? std::make_optional(raft::device_span<vertex_t const>{
+                      start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                  : std::nullopt,
+                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                      : std::nullopt,
+                offsets ? std::make_optional(
+                            raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                        : std::nullopt,
+                edge_label ? edge_label->size() : size_t{1},
+                hop ? fan_out_->size_ : size_t{1},
+                src_is_major,
+                do_expensive_check_);
+
+            majors.emplace(std::move(output_majors));
+            renumber_map.emplace(std::move(output_renumber_map));
+          } else {
+            // (D)CSC, (D)CSR
+
+            bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+                                  (options_.compression_type_ == cugraph_compression_type_t::DCSC);
+
+            rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
+            rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+            std::tie(majors,
+                    output_major_offsets,
+                    minors,
+                    wgt,
+                    edge_id,
+                    edge_type,
+                    label_hop_offsets,
+                    output_renumber_map,
+                    renumber_map_offsets) =
+              cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
+                handle_,
+                std::move(src),
+                std::move(dst),
+                std::move(wgt),
+                std::move(edge_id),
+                std::move(edge_type),
+                std::move(hop),
+                options_.retain_seeds_
+                  ? std::make_optional(raft::device_span<vertex_t const>{
+                      start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                  : std::nullopt,
+                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                      : std::nullopt,
+                offsets ? std::make_optional(
+                            raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                        : std::nullopt,
+                edge_label ? edge_label->size() : size_t{1},
+                hop ? fan_out_->size_ : size_t{1},
+                src_is_major,
+                options_.compress_per_hop_,
+                doubly_compress,
+                do_expensive_check_);
+
+            renumber_map.emplace(std::move(output_renumber_map));
+            major_offsets.emplace(std::move(output_major_offsets));
+          }
+
+          // These are now represented by label_hop_offsets
+          hop.reset();
+          offsets.reset();
+        
+        } else { // heterogeneous
 
           rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
-          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-          std::tie(output_majors,
-                   minors,
-                   wgt,
-                   edge_id,
-                   edge_type,
-                   label_hop_offsets,
-                   output_renumber_map,
-                   renumber_map_offsets) =
-            cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
-              handle_,
-              std::move(src),
-              std::move(dst),
-              std::move(wgt),
-              std::move(edge_id),
-              std::move(edge_type),
-              std::move(hop),
-              options_.retain_seeds_
-                ? std::make_optional(raft::device_span<vertex_t const>{
-                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
-                : std::nullopt,
-              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
-                                     : std::nullopt,
-              offsets ? std::make_optional(
-                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
-                      : std::nullopt,
-              edge_label ? edge_label->size() : size_t{1},
-              hop ? fan_out_->size_ : size_t{1},
-              src_is_major,
-              do_expensive_check_);
-
-          majors.emplace(std::move(output_majors));
-          renumber_map.emplace(std::move(output_renumber_map));
-        } else {
-          // (D)CSC, (D)CSR
-
-          bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
-                                 (options_.compression_type_ == cugraph_compression_type_t::DCSC);
-
-          rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
-          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-          std::tie(majors,
-                   output_major_offsets,
-                   minors,
-                   wgt,
-                   edge_id,
-                   edge_type,
-                   label_hop_offsets,
-                   output_renumber_map,
-                   renumber_map_offsets) =
-            cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
-              handle_,
-              std::move(src),
-              std::move(dst),
-              std::move(wgt),
-              std::move(edge_id),
-              std::move(edge_type),
-              std::move(hop),
-              options_.retain_seeds_
-                ? std::make_optional(raft::device_span<vertex_t const>{
-                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
-                : std::nullopt,
-              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
-                                     : std::nullopt,
-              offsets ? std::make_optional(
-                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
-                      : std::nullopt,
-              edge_label ? edge_label->size() : size_t{1},
-              hop ? fan_out_->size_ : size_t{1},
-              src_is_major,
-              options_.compress_per_hop_,
-              doubly_compress,
-              do_expensive_check_);
-
-          renumber_map.emplace(std::move(output_renumber_map));
-          major_offsets.emplace(std::move(output_major_offsets));
+            rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+            std::tie(output_majors,
+                    minors,
+                    wgt,
+                    edge_id,
+                    edge_type,
+                    label_hop_offsets,
+                    output_renumber_map,
+                    renumber_map_offsets) =
+              cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
+                handle_,
+                std::move(src),
+                std::move(dst),
+                std::move(wgt),
+                std::move(edge_id),
+                std::move(edge_type),
+                std::move(hop),
+                options_.retain_seeds_
+                  ? std::make_optional(raft::device_span<vertex_t const>{
+                      start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                  : std::nullopt,
+                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                      : std::nullopt,
+                offsets ? std::make_optional(
+                            raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                        : std::nullopt,
+                edge_label ? edge_label->size() : size_t{1},
+                hop ? fan_out_->size_ : size_t{1},
+                src_is_major,
+                do_expensive_check_);
+
+            majors.emplace(std::move(output_majors));
+            renumber_map.emplace(std::move(output_renumber_map));
+        
         }
 
-        // These are now represented by label_hop_offsets
-        hop.reset();
-        offsets.reset();
+
+
+
+
       } else {
         if (options_.compression_type_ != cugraph_compression_type_t::COO) {
           CUGRAPH_FAIL("Can only use COO format if not renumbering");

From d7165af2eeecc4fc2c148203e2a147fbf38101d5 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 17 Sep 2024 06:49:07 -0700
Subject: [PATCH 035/117] update heterogeneous renumbering call

---
 cpp/src/c_api/neighbor_sampling.cpp | 111 ++++++++++++++--------------
 1 file changed, 54 insertions(+), 57 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 6348c6f9731..afe629501d7 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -415,7 +415,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           offsets.reset();
         
         } else { // heterogeneous
-          //graph_view.local_vertex_partition_range_size()
 
           rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
 
@@ -424,69 +423,67 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                                 vertex_type_offsets.size(),
                                 vertex_t{0});
 
-          //std::get<0>(*heterogeneous_fan_out_).back() - 1
-
           rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
-            rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-            /*
-            std::tie(output_majors,
-                    minors,
-                    wgt,
-                    edge_id,
-                    edge_type,
-                    label_hop_offsets,
-                    output_renumber_map,
-                    renumber_map_offsets) =
-            */
-            /*
-            std::tie(output_majors,
-               minors,
-               wgt,
-               edge_id,
-               edge_type,
-               renumbered_and_sorted_vertex_renumber_map,
-               renumbered_and_sorted_vertex_renumber_map_label_type_offsets,
-               renumbered_and_sorted_edge_id_renumber_map,
-               renumbered_and_sorted_edge_id_renumber_map_label_type_offsets) =
-            */
-            auto x =
-              cugraph::heterogeneous_renumber_and_sort_sampled_edgelist<vertex_t>(
-                handle_,
-                std::move(src),
-                std::move(dst),
-                std::move(wgt),
-                std::move(edge_id),
-                std::move(edge_type),
-                std::move(hop),
-                options_.retain_seeds_
-                  ? std::make_optional(raft::device_span<vertex_t const>{
-                      start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
-                  : std::nullopt,
-                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
-                                      : std::nullopt,
-                offsets ? std::make_optional(
-                            raft::device_span<size_t const>{offsets->data(), offsets->size()})
-                        : std::nullopt,
-                raft::device_span<vertex_t const>{
-                      vertex_type_offsets.data(), vertex_type_offsets.size()},
-                
-                edge_label ? edge_label->size() : size_t{1},
-                hop ? fan_out_->size_ : size_t{1},
-                size_t{1},
-                size_t{1},
-                src_is_major,
-                do_expensive_check_);
+          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+          
+          std::optional<rmm::device_uvector<edge_t>> renumbered_and_sorted_edge_id_renumber_map(std::nullopt);
+          std::optional<rmm::device_uvector<size_t>>    renumbered_and_sorted_edge_id_renumber_map_label_type_offsets(std::nullopt);
+
+          // extract the edge_type from label_type_hop_offsets
+          std::optional<rmm::device_uvector<size_t>> label_type_hop_offsets{std::nullopt};
+
+          std::tie(output_majors,
+              minors,
+              wgt,
+              edge_id,
+              label_type_hop_offsets, // Contains information about the type and hop offsets
+              output_renumber_map,
+              (*renumber_map_offsets),
+              renumbered_and_sorted_edge_id_renumber_map,
+              renumbered_and_sorted_edge_id_renumber_map_label_type_offsets
+              ) = 
+            cugraph::heterogeneous_renumber_and_sort_sampled_edgelist<vertex_t>(
+              handle_,
+              std::move(src),
+              std::move(dst),
+              std::move(wgt),
+              std::move(edge_id),
+              std::move(edge_type),
+              std::move(hop),
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                        label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                    : std::nullopt,
+              offsets ? std::make_optional(
+                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                      : std::nullopt,
+              raft::device_span<vertex_t const>{
+                    vertex_type_offsets.data(), vertex_type_offsets.size()},
+              
+              edge_label ? edge_label->size() : size_t{1},
+              hop ? fan_out_->size_ : size_t{1},
+              size_t{1},
+              size_t{1},
+              src_is_major,
+              do_expensive_check_);
+            if (edge_type) {
+              (*edge_type).resize(raft::device_span<size_t const>{
+                (*label_type_hop_offsets).data(), (*label_type_hop_offsets).size()}.back() - 1, handle_.get_stream());
+              cugraph::detail::sequence_fill(handle_.get_stream(),
+                                (*edge_type).begin(),
+                                (*edge_type).size(),
+                                edge_type_t{0});
+            }
+            
 
             majors.emplace(std::move(output_majors));
             renumber_map.emplace(std::move(output_renumber_map));
         
         }
 
-
-
-
-
       } else {
         if (options_.compression_type_ != cugraph_compression_type_t::COO) {
           CUGRAPH_FAIL("Can only use COO format if not renumbering");

From 5cdf40ab7c243590d70bfddf6de257669fa0e994 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 17 Sep 2024 09:16:02 -0700
Subject: [PATCH 036/117] update API and docstring

---
 cpp/include/cugraph/sampling_functions.hpp  | 187 ++-------
 cpp/src/sampling/neighbor_sampling_impl.hpp | 399 +-------------------
 2 files changed, 43 insertions(+), 543 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 29904450562..35dbf08a6c9 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -178,6 +178,7 @@ uniform_neighbor_sample(
  * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
  * @tparam weight_t Type of edge weights. Needs to be a floating point type.
  * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam bias_t Type of bias. Needs to be an integral type.
  * @tparam label_t Type of label. Needs to be an integral type.
  * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
  * true) are major indices
@@ -250,80 +251,6 @@ biased_neighbor_sample(
   bool dedupe_sources                             = false,
   bool do_expensive_check                         = false);
 
-/**
- * @brief Heterogeneous Neighborhood Sampling.
- *
- * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
- * The branching out to select outgoing neighbors is performed with homogeneous fanouts.
- *
- * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
- * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
- * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
- * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
- * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
- * (optional) will be described below and is dependent upon the input parameters.
- *
- * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
- * label and offsets values in the return set will be std::nullopt.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
- * the label output has values.  This will also result in the output being sorted by vertex label.
- * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
- * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
- * label output has values.  This will also result in the output being sorted by vertex label.  The
- * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
- * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
- * be shuffled so that all data with a particular label will be on the specified rank.
- *
- * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
- * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
- * @tparam weight_t Type of edge weights. Needs to be a floating point type.
- * @tparam edge_type_t Type of edge type. Needs to be an integral type.
- * @tparam label_t Type of label. Needs to be an integral type.
- * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
- * true) are major indices
- * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
- * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
- * handles to various CUDA libraries) to run graph algorithms.
- * * @param rng_state A pre-initialized raft::RngState object for generating random numbers
- * @param graph_view Graph View object to generate NBR Sampling on.
- * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
- * @param edge_id_view Optional view object holding edge ids for @p graph_view.
- * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
- * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
- * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
- * uniform sampling.
- * @param starting_vertices Device span of starting vertex IDs for the sampling.
- * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
- * for the sampling.
- * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
- * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
- * output rank.  The label span must be sorted in ascending order.
- * @param heterogeneous_fan_out Tuple of host spans defining branching out (fan-out) degree per
- * source vertex for each level in CSR style format. The first element of the tuple is the offset
- * array per edge type id and the second element corresponds to the fanout values.
- * The sampling method can use different fan_out values for each edge type.
- * The fan-out offsets size must be proportional to the number of edge types and fan_out values.
- * @param return_hops boolean flag specifying if the hop information should be returned. 
- * @param prior_sources_behavior Enum type defining how to handle prior sources, (defaults to
- * DEFAULT)
- * @param dedupe_sources boolean flag, if true then if a vertex v appears as a destination in hop X
- * multiple times with the same label, it will only be passed once (for each label) as a source
- * for the next hop.  Default is false.
- * @param with_replacement boolean flag specifying if random sampling is done with replacement
- * (true); or, without replacement (false); default = true;
- * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
- * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
- * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
- * optional int32_t hop, optional label_t label, optional size_t offsets)
- */
-
 struct sampling_flags_t {
   /**
    * Specifies how to handle prior sources. Default is DEFAULT.
@@ -349,77 +276,12 @@ struct sampling_flags_t {
   bool with_replacement{true};
 };
 
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename bias_t,
-          typename label_t,
-          bool store_transposed,
-          bool multi_gpu>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>,
-           std::optional<rmm::device_uvector<edge_type_t>>,
-           std::optional<rmm::device_uvector<int32_t>>,
-           std::optional<rmm::device_uvector<label_t>>,
-           std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
-  raft::device_span<vertex_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement                           = true,
-  prior_sources_behavior_t prior_sources_behavior = prior_sources_behavior_t::DEFAULT,
-  bool dedupe_sources                             = false,
-  bool do_expensive_check                         = false);
-
 /**
- * @brief Homogeneous Neighborhood Sampling.
+ * @brief Heterogeneous Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
  * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
- * The branching out to select outgoing neighbors is performed with homogeneous fanouts
- *
- * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
- * label, offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
- * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
- * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
- * encountered in.  The label output (optional) identifes the vertex label.  The offsets array
- * (optional) will be described below and is dependent upon the input parameters.
- *
- * If @p starting_vertex_labels is not specified then no organization is applied to the output, the
- * label and offsets values in the return set will be std::nullopt.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is not specified then
- * the label output has values.  This will also result in the output being sorted by vertex label.
- * The offsets array in the return will be a CSR-style offsets array to identify the beginning of
- * each label range in the data.  `labels.size() == (offsets.size() - 1)`.
- *
- * If @p starting_vertex_labels is specified and @p label_to_output_comm_rank is specified then the
- * label output has values.  This will also result in the output being sorted by vertex label.  The
- * offsets array in the return will be a CSR-style offsets array to identify the beginning of each
- * label range in the data.  `labels.size() == (offsets.size() - 1)`.  Additionally, the data will
- * be shuffled so that all data with a particular label will be on the specified rank.
- * @brief Defines flags for sampling
- */
-
-/**
- * @brief Uniform Neighborhood Sampling.
- *
- * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with homogeneous fanouts.
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
  * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
@@ -442,6 +304,7 @@ heterogeneous_neighbor_sample(
  * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
  * @tparam weight_t Type of edge weights. Needs to be a floating point type.
  * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam bias_t Type of bias. Needs to be an integral type.
  * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
  * true) are major indices
  * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
@@ -452,6 +315,11 @@ heterogeneous_neighbor_sample(
  * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
  * @param edge_id_view Optional view object holding edge ids for @p graph_view.
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
+ * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
+ * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
+ * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
+ * uniform sampling.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
@@ -460,16 +328,17 @@ heterogeneous_neighbor_sample(
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
  * level
- * @param flags A set of flags indicating which sampling features should be used.l
+ * @param flags A set of flags indicating which sampling features should be used.
  * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
  * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
  * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
- * optional int32_t hop, optional size_t offsets)
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
  */
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
           typename edge_type_t,
+          typename bias_t,
           bool store_transposed,
           bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>,
@@ -479,25 +348,28 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
   sampling_flags_t sampling_flags,
   bool do_expensive_check = false);
 
 /**
- * @brief Biased Neighborhood Sampling.
+ * @brief Homogeneous Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
+ * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with homogeneous fanouts
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
  * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
@@ -512,10 +384,15 @@ homogeneous_uniform_neighbor_sample(
  * in the return will be a CSR-style offsets array to identify the beginning of each label range in
  * the output vectors.
  *
+ * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
+ * for a particular label are returned on the specified rank.  This will result in the offsets array
+ * on other GPUs indicating that there are no entries for that label (`offsets[i] == offsets[i+1]`).
+ *
  * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
  * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
  * @tparam weight_t Type of edge weights. Needs to be a floating point type.
  * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam bias_t Type of bias. Needs to be an integral type.
  * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
  * true) are major indices
  * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
@@ -526,24 +403,26 @@ homogeneous_uniform_neighbor_sample(
  * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
  * @param edge_id_view Optional view object holding edge ids for @p graph_view.
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
+ * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
  * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
  * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected.
+ * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
+ * uniform sampling.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_offsets Optional device span of offsets identifying the range of
+  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
  * starting vertex values for this label.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
  * level
- * @param flags A set of flags indicating which sampling features should be used.l
+ * @param flags A set of flags indicating which sampling features should be used.
  * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
  * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
  * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
- * optional int32_t hop, optional size_t offsets)
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
  */
+
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
@@ -558,14 +437,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 8f220272f4a..eb09053feec 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -552,6 +552,7 @@ template <typename vertex_t,
           typename edge_t,
           typename weight_t,
           typename edge_type_t,
+          typename bias_t,
           bool store_transposed,
           bool multi_gpu>
 std::tuple<rmm::device_uvector<vertex_t>,
@@ -561,22 +562,24 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
+heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
+  std::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
+    heterogeneous_fan_out,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
   using bias_t = weight_t;  // dummy
-
+  
   rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
 
   if (starting_vertex_offsets)
@@ -590,7 +593,7 @@ homogeneous_uniform_neighbor_sample(
     edge_weight_view,
     edge_id_view,
     edge_type_view,
-    std::nullopt,
+    edge_bias_view,
     starting_vertices,
     starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
                                 starting_vertex_labels.data(), starting_vertex_labels.size()})
@@ -618,14 +621,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
+homogeneous_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
+  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -633,6 +636,7 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
+
   rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
 
   if (starting_vertex_offsets)
@@ -641,7 +645,6 @@ homogeneous_biased_neighbor_sample(
 
   return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
-    rng_state,
     graph_view,
     edge_weight_view,
     edge_id_view,
@@ -653,378 +656,6 @@ homogeneous_biased_neighbor_sample(
                             : std::nullopt,
     label_to_output_comm_rank,
     fan_out,
-    sampling_flags.return_hops,
-    sampling_flags.with_replacement,
-    sampling_flags.prior_sources_behavior,
-    sampling_flags.dedupe_sources,
-    do_expensive_check);
-}
-
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename bias_t,
-          typename label_t,
-          bool store_transposed,
-          bool multi_gpu>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>,
-           std::optional<rmm::device_uvector<edge_type_t>>,
-           std::optional<rmm::device_uvector<int32_t>>,
-           std::optional<rmm::device_uvector<label_t>>,
-           std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
-  raft::device_span<vertex_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check)
-{
-  // for loop : over the number of edgetypes /
-  // inside the loop, create an edge mask and update the edge mask to mask out all the irrelevant edges
-  // adn call neighborhood sampling
-
-  // before this, there is a reordering constrain.
-  // 10 labels amd 100 seeds, 3 edge type. For each seed, you provide sample of each edge type (0, 1, 2)
-  // soritng criterial, label , seed, edge tyope
-  // label, edge type, seed
-
-  // Pick how to sort - requries assembling? if the natural output order based on this scheme if it doesn't meet the requirement, needs to shuffle
-  // shuffling, label to be the primiary key? if you loop over the edge tyoe, edge type will be the primary key.
-  
-  // Goal: reorder the sample at the end - Label is the primary key? then label? then seed?
-  // stable sort? to not mess up the internal order? "Use scatter operation to be more efficient?"
-
-  // Iterate over all edge types
-
-  cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
-  cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
-
-
-  std::vector<rmm::device_uvector<vertex_t>> level_result_src_vectors{};
-  std::vector<rmm::device_uvector<vertex_t>> level_result_dst_vectors{};
-  auto level_result_weight_vectors =
-    edge_weight_view ? std::make_optional(std::vector<rmm::device_uvector<weight_t>>{})
-                     : std::nullopt;
-  auto level_result_edge_id_vectors =
-    edge_id_view ? std::make_optional(std::vector<rmm::device_uvector<edge_t>>{}) : std::nullopt;
-  auto level_result_edge_type_vectors =
-    edge_type_view ? std::make_optional(std::vector<rmm::device_uvector<edge_type_t>>{})
-                   : std::nullopt;
-  auto level_result_label_vectors =
-    starting_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
-                                : std::nullopt;
-  
-    level_result_src_vectors.reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
-  level_result_dst_vectors.reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
-  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
-  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
-  if (level_result_edge_type_vectors) {
-    (*level_result_edge_type_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1);
-  }
-  if (level_result_label_vectors) { (*level_result_label_vectors).reserve(std::get<0>(heterogeneous_fan_out).size() - 1); }
-
-
-  
-
-
-
-
-
-
-
-
-  std::vector<size_t> level_sizes{};
-  for (int i = 0; i < std::get<0>(heterogeneous_fan_out).size() - 1; i++) {
-
-    cugraph::transform_e(
-        handle,
-        graph_view,
-        cugraph::edge_src_dummy_property_t{}.view(),
-        cugraph::edge_dst_dummy_property_t{}.view(),
-        *edge_type_view,
-        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
-          return edge_type == valid_edge_type;
-        },
-        edge_mask.mutable_view(),
-        false);
-    
-    auto [
-      level_result_srcs,
-      level_result_dsts,
-      level_result_weights,
-      level_result_edge_ids,
-      level_result_edge_types,
-      level_result_hops,
-      level_result_labels,
-      label_to_output_comm_rank_ // unused
-      ] = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-        handle,
-        graph_view,
-        edge_weight_view,
-        edge_id_view,
-        edge_type_view,
-        edge_bias_view,
-        starting_vertices,
-        starting_vertex_labels,
-        label_to_output_comm_rank,
-        std::nullopt,
-        heterogeneous_fan_out,
-        return_hops,
-        with_replacement,
-        prior_sources_behavior,
-        dedupe_sources,
-        rng_state,
-        do_expensive_check);
-    
-    level_sizes.push_back(level_result_srcs.size());
-
-    level_result_src_vectors.push_back(std::move(level_result_srcs));
-    level_result_dst_vectors.push_back(std::move(level_result_dsts));
-    if (level_result_weights) { (*level_result_weight_vectors).push_back(std::move(*level_result_weights)); }
-    if (level_result_edge_ids) { (*level_result_edge_id_vectors).push_back(std::move(*level_result_edge_ids)); }
-    if (level_result_edge_types) { (*level_result_edge_type_vectors).push_back(std::move(*level_result_edge_types)); }
-    if (level_result_labels) { (*level_result_label_vectors).push_back(std::move(*level_result_labels)); }
-  }
-
-
-  auto result_size = std::reduce(level_sizes.begin(), level_sizes.end());
-
-  size_t output_offset{};
-
-  rmm::device_uvector<vertex_t> result_srcs(result_size, handle.get_stream());
-  output_offset = 0;
-  for (size_t i = 0; i < level_result_src_vectors.size(); ++i) {
-    raft::copy(result_srcs.begin() + output_offset,
-               level_result_src_vectors[i].begin(),
-               level_sizes[i],
-               handle.get_stream());
-    output_offset += level_sizes[i];
-  }
-  level_result_src_vectors.clear();
-  level_result_src_vectors.shrink_to_fit();
-
-  rmm::device_uvector<vertex_t> result_dsts(result_size, handle.get_stream());
-  output_offset = 0;
-  for (size_t i = 0; i < level_result_dst_vectors.size(); ++i) {
-    raft::copy(result_dsts.begin() + output_offset,
-               level_result_dst_vectors[i].begin(),
-               level_sizes[i],
-               handle.get_stream());
-    output_offset += level_sizes[i];
-  }
-  level_result_dst_vectors.clear();
-  level_result_dst_vectors.shrink_to_fit();
-
-  auto result_weights =
-    level_result_weight_vectors
-      ? std::make_optional(rmm::device_uvector<weight_t>(result_size, handle.get_stream()))
-      : std::nullopt;
-  if (result_weights) {
-    output_offset = 0;
-    for (size_t i = 0; i < (*level_result_weight_vectors).size(); ++i) {
-      raft::copy((*result_weights).begin() + output_offset,
-                 (*level_result_weight_vectors)[i].begin(),
-                 level_sizes[i],
-                 handle.get_stream());
-      output_offset += level_sizes[i];
-    }
-    level_result_weight_vectors = std::nullopt;
-  }
-
-  auto result_edge_ids =
-    level_result_edge_id_vectors
-      ? std::make_optional(rmm::device_uvector<edge_t>(result_size, handle.get_stream()))
-      : std::nullopt;
-  if (result_edge_ids) {
-    output_offset = 0;
-    for (size_t i = 0; i < (*level_result_edge_id_vectors).size(); ++i) {
-      raft::copy((*result_edge_ids).begin() + output_offset,
-                 (*level_result_edge_id_vectors)[i].begin(),
-                 level_sizes[i],
-                 handle.get_stream());
-      output_offset += level_sizes[i];
-    }
-    level_result_edge_id_vectors = std::nullopt;
-  }
-
-  auto result_edge_types =
-    level_result_edge_type_vectors
-      ? std::make_optional(rmm::device_uvector<edge_type_t>(result_size, handle.get_stream()))
-      : std::nullopt;
-  if (result_edge_types) {
-    output_offset = 0;
-    for (size_t i = 0; i < (*level_result_edge_type_vectors).size(); ++i) {
-      raft::copy((*result_edge_types).begin() + output_offset,
-                 (*level_result_edge_type_vectors)[i].begin(),
-                 level_sizes[i],
-                 handle.get_stream());
-      output_offset += level_sizes[i];
-    }
-    level_result_edge_type_vectors = std::nullopt;
-  }
-
-  std::optional<rmm::device_uvector<int32_t>> result_hops{std::nullopt};
-
-  if (return_hops) {
-    result_hops   = rmm::device_uvector<int32_t>(result_size, handle.get_stream());
-    output_offset = 0;
-    for (size_t i = 0; i < std::get<0>(heterogeneous_fan_out).size() - 1; ++i) {
-      detail::scalar_fill(
-        handle, result_hops->data() + output_offset, level_sizes[i], static_cast<int32_t>(i));
-      output_offset += level_sizes[i];
-    }
-  }
-
-  auto result_labels =
-    level_result_label_vectors
-      ? std::make_optional(rmm::device_uvector<label_t>(result_size, handle.get_stream()))
-      : std::nullopt;
-  if (result_labels) {
-    output_offset = 0;
-    for (size_t i = 0; i < (*level_result_label_vectors).size(); ++i) {
-      raft::copy((*result_labels).begin() + output_offset,
-                 (*level_result_label_vectors)[i].begin(),
-                 level_sizes[i],
-                 handle.get_stream());
-      output_offset += level_sizes[i];
-    }
-    level_result_label_vectors = std::nullopt;
-  }
-
-  // 
-
-  /*
-  return std::make_tuple(
-    std::move(result_srcs),
-    std::move(result_dsts),
-    std::move(result_weights),
-    std::move(result_edge_ids),
-    std::move(result_edge_types),
-    std::move(result_hops),
-    std::move(result_labels),
-    //offsets
-    //std::move(label_to_output_comm_rank)
-  );
-  */
-
-  // Remove and pick the one on top
-  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    edge_bias_view,
-    starting_vertices,
-    starting_vertex_labels,
-    label_to_output_comm_rank,
-    std::nullopt,
-    heterogeneous_fan_out,
-    return_hops,
-    with_replacement,
-    prior_sources_behavior,
-    dedupe_sources,
-    rng_state,
-    do_expensive_check);
-
-
-
-
-
-
-
-}
-
-  /*
-  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    edge_bias_view,
-    starting_vertices,
-    starting_vertex_labels,
-    label_to_output_comm_rank,
-    std::nullopt,
-    heterogeneous_fan_out,
-    return_hops,
-    with_replacement,
-    prior_sources_behavior,
-    dedupe_sources,
-    rng_state,
-    do_expensive_check);
-//}
-*/
-
-
-
-
-
-
-template <typename vertex_t,
-          typename edge_t,
-          typename weight_t,
-          typename edge_type_t,
-          typename bias_t,
-          typename label_t,
-          bool store_transposed,
-          bool multi_gpu>
-std::tuple<rmm::device_uvector<vertex_t>,
-           rmm::device_uvector<vertex_t>,
-           std::optional<rmm::device_uvector<weight_t>>,
-           std::optional<rmm::device_uvector<edge_t>>,
-           std::optional<rmm::device_uvector<edge_type_t>>,
-           std::optional<rmm::device_uvector<int32_t>>,
-           std::optional<rmm::device_uvector<label_t>>,
-           std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
-  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
-  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
-  raft::device_span<vertex_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check)
-{
-  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    edge_bias_view,
-    starting_vertices,
-    starting_vertex_labels,
-    label_to_output_comm_rank,
-    fan_out,
     std::nullopt,
     return_hops,
     with_replacement,
@@ -1034,14 +665,4 @@ homogeneous_neighbor_sample(
     do_expensive_check);
 }
 
-
-
-
-
-
-
-
-
-
-
 }  // namespace cugraph

From a8fbd9d19e401c0a699b57ad95f5b841573e19c2 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 17 Sep 2024 09:16:58 -0700
Subject: [PATCH 037/117] remove unsued variable

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index eb09053feec..6b27b87c7f5 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -577,9 +577,7 @@ heterogeneous_neighbor_sample(
     heterogeneous_fan_out,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
-{
-  using bias_t = weight_t;  // dummy
-  
+{ 
   rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
 
   if (starting_vertex_offsets)

From 9d5b3dd05c9f5afefb86c3cff5bba87590330225 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 09:35:58 -0700
Subject: [PATCH 038/117] update C++ API for neighbor sampling

---
 cpp/include/cugraph/sampling_functions.hpp    |  10 +-
 cpp/include/cugraph_c/sampling_algorithms.h   |  67 ++--
 cpp/src/c_api/neighbor_sampling.cpp           | 322 +++++++++++-------
 cpp/src/sampling/neighbor_sampling_impl.hpp   | 284 +++++++--------
 .../sampling/neighbor_sampling_mg_v32_e32.cu  |  52 +--
 .../sampling/neighbor_sampling_mg_v32_e64.cu  |  52 +--
 .../sampling/neighbor_sampling_mg_v64_e64.cu  |  52 +--
 .../sampling/neighbor_sampling_sg_v32_e32.cu  |  52 +--
 .../sampling/neighbor_sampling_sg_v32_e64.cu  |  52 +--
 .../sampling/neighbor_sampling_sg_v64_e64.cu  |  52 +--
 10 files changed, 474 insertions(+), 521 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 35dbf08a6c9..c59791d4a62 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -327,7 +327,9 @@ struct sampling_flags_t {
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level
+ * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types + edge_type_id]'
+ * @param num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
+ * sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param flags A set of flags indicating which sampling features should be used.
  * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
  * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
@@ -359,8 +361,8 @@ heterogeneous_neighbor_sample(
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  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,
+  edge_type_t num_edge_types,
   sampling_flags_t sampling_flags,
   bool do_expensive_check = false);
 
@@ -415,7 +417,7 @@ heterogeneous_neighbor_sample(
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level
+ * level.
  * @param flags A set of flags indicating which sampling features should be used.
  * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
  * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index ba6e33b7575..b9e7836b550 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -199,6 +199,13 @@ typedef struct {
   int32_t align_;
 } cugraph_sampling_options_t;
 
+/**
+ * @brief     Opaque sampling options type
+ */
+typedef struct {
+  int32_t align_;
+} sampling_flags_t;
+
 /**
  * @brief     Enumeration for prior sources behavior
  */
@@ -321,29 +328,29 @@ void cugraph_sampling_set_dedupe_sources(cugraph_sampling_options_t* options, bo
 void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
 
 /**
- * @brief     Opaque neighborhood sampling heterogeneous fanout type
+ * @brief     Opaque neighborhood sampling heterogeneous fan_out type
  */
 typedef struct {
   int32_t align_;
 } cugraph_sample_heterogeneous_fan_out_t;
 
 /**
- * @brief       Create heterogeneous fanout
+ * @brief       Create heterogeneous fan_out
  *
- * Input data will be stored in the heterogenous_fanout.
+ * Input data will be stored in the heterogenous_fan_out.
  * 
- * The fanout is going to be a CSR structure, the edge_type_offsets will define which range
- * of the fanout array is associated with each edge type, the fanout will be the values of
- * fanout for that hop/type. So for edge type k, fanout[edge_type_offsets[k]] will identify
- * the fanout for hop 0 for edge type k.  fanout[edge_type_offsets[k] +1] will identify the
- * fanout for hop 1, etc. edge_type_offsets[k+1] will mark the beginning of the fanout
- * array for type k+1 (and the end of the fanout array for type k.
+ * The fan_out is going to be a CSR structure, the edge_type_offsets will define which range
+ * of the fan_out array is associated with each edge type, the fan_out will be the values of
+ * fan_out for that hop/type. So for edge type k, fan_out[edge_type_offsets[k]] will identify
+ * the fan_out for hop 0 for edge type k.  fan_out[edge_type_offsets[k] +1] will identify the
+ * fan_out for hop 1, etc. edge_type_offsets[k+1] will mark the beginning of the fan_out
+ * array for type k+1 (and the end of the fan_out array for type k.
  *
  * @param [in]  handle         Handle for accessing resources
  * @param [in]  graph          Pointer to graph
  * @param [in]  edge_type_offsets Type erased array of edge type offsets
- * @param [in]  fanout         Type erased array of fanout values
- * @param [out] heterogeneous_fanout Opaque pointer to fanout_t
+ * @param [in]  fan_out         Type erased array of fan_out values
+ * @param [out] heterogeneous_fan_out Opaque pointer to fan_out_t
  * @param [out] error          Pointer to an error object storing details of any error.  Will
  *                             be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
@@ -352,16 +359,16 @@ cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
   const cugraph_type_erased_host_array_view_t* edge_type_offsets,
-  const cugraph_type_erased_host_array_view_t* fanout,
-  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fanout,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
   cugraph_error_t** error);
 
 /**
- * @brief     Free edge type and fanout pairs
+ * @brief     Free edge type and fan_out pairs
  *
- * @param [in]    heterogeneous_fanout The edge type size and fanout values
+ * @param [in]    heterogeneous_fan_out The edge type size and fan_out values
  */
-void cugraph_heterogeneous_fan_out_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
+void cugraph_heterogeneous_fan_out_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out);
 
 /**
  * @brief     Uniform Neighborhood Sampling
@@ -394,8 +401,8 @@ void cugraph_heterogeneous_fan_out_free(cugraph_sample_heterogeneous_fan_out_t*
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fanout values of type INT32
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                           We only support fan_out values of type INT32
  * @param [in,out] rng_state State of the random number generator, updated with each call
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
@@ -455,8 +462,8 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fanout values of type INT32
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                           We only support fan_out values of type INT32
  * @param [in,out] rng_state State of the random number generator, updated with each call
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
@@ -487,7 +494,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * @brief     Heterogeneous Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices with edge biases or not
- * and homogeneous fanout types.
+ * and homogeneous fan_out types.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -516,10 +523,11 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  heterogeneous_fanout  Tuple of host arrays defining the fan out at each step in the
- * sampling algorithm. in CSR style format. The first element of the tuple is the offset array per
- * edge type id and the second element corresponds to the fanout values.
- * We only support type INT32 for both the offsets and the fanout values array.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                            We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
+ *                            sample whereas a value greater than 1 translates to heterogeneous neighbor
+ *                            sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  is_biased
@@ -542,7 +550,8 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
-  const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
+  const cugraph_sample_heterogeneous_fan_out_t* fan_out,
+  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
@@ -553,7 +562,7 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
  * @brief     Homogeneous Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices with edge biases or not
- * and homogeneous fanout types.
+ * and homogeneous fan_out types.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -582,8 +591,8 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fanout       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fanout values of type INT32
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                            We only support fan_out values of type INT32
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  is_biased
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index afe629501d7..62f55c19472 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -47,6 +47,13 @@ struct cugraph_sampling_options_t {
   bool_t retain_seeds_{FALSE};
 };
 
+struct sampling_flags_t {
+  prior_sources_behavior_t prior_sources_behavior_{prior_sources_behavior_t::DEFAULT};
+  bool_t return_hops_{FALSE};
+  bool_t dedupe_sources_{FALSE};
+  bool_t with_replacement_{FALSE};
+};
+
 struct cugraph_sample_result_t {
   cugraph_type_erased_device_array_t* major_offsets_{nullptr};
   cugraph_type_erased_device_array_t* majors_{nullptr};
@@ -73,13 +80,15 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
-  cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const* heterogeneous_fan_out_{nullptr};
+  int num_edge_types_{};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool is_biased_{false};
+  bool is_deprecated_{false};
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
@@ -90,13 +99,15 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
     cugraph_type_erased_device_array_view_t const* start_vertex_labels,
+    cugraph_type_erased_device_array_view_t const* start_vertex_offsets,
     cugraph_type_erased_device_array_view_t const* label_list,
     cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
     cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
-    cugraph_sample_heterogeneous_fan_out_t const* heterogeneous_fan_out,
+    int num_edge_types,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool is_biased,
+    bool is_deprecated,
     bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
@@ -110,6 +121,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       start_vertex_labels_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           start_vertex_labels)),
+      start_vertex_offsets_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertex_offsets)),
       label_list_(reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
         label_list)),
       label_to_comm_rank_(
@@ -120,13 +134,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           label_offsets)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
-
-      heterogeneous_fan_out_(
-        reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(
-          heterogeneous_fan_out)),
-  
+      num_edge_types_(num_edge_types),
       options_(options),
       is_biased_(is_biased),
+      is_deprecated_(is_deprecated),
       do_expensive_check_(do_expensive_check)
   {
   }
@@ -194,6 +205,17 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                    handle_.get_stream());
       }
 
+      std::optional<rmm::device_uvector<size_t>> start_vertex_offsets{std::nullopt};
+
+      if (start_vertex_offsets_ != nullptr) {
+        start_vertex_offsets =
+          rmm::device_uvector<size_t>{start_vertex_offsets_->size_, handle_.get_stream()};
+        raft::copy(start_vertex_offsets->data(),
+                   start_vertex_offsets_->as_type<size_t>(),
+                   start_vertex_offsets_->size_,
+                   handle_.get_stream());
+      }
+
       if constexpr (multi_gpu) {
         if (start_vertex_labels) {
           std::tie(start_vertices, *start_vertex_labels) =
@@ -226,73 +248,142 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       std::optional<rmm::device_uvector<int32_t>> hop{std::nullopt};
       std::optional<rmm::device_uvector<label_t>> edge_label{std::nullopt};
       std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
-      // FIXME: Consolidate 'fan_out_' and 'heterogeneous_fan_out_' into one
-      // argument with std::variant
-      //auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-      if (heterogeneous_fan_out_ != nullptr) {
-        std::tie(src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets) =
-          cugraph::heterogeneous_neighbor_sample(
-            handle_,
-            rng_state_->rng_state_,
-            graph_view,
-            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-            (start_vertex_labels_ != nullptr)
-              ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                    start_vertex_labels->size())
-              : std::nullopt,
-            (label_list_ != nullptr)
-              ? std::make_optional(std::make_tuple(
-                  raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                  label_list_->size_},
-                  raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                  label_to_comm_rank_->size_}))
-              : std::nullopt,
-
-            std::make_tuple(
-              raft::host_span<const int>{std::get<0>(*heterogeneous_fan_out_)->as_type<int>(),
-                                        std::get<0>(*heterogeneous_fan_out_)->size_},
-              raft::host_span<const int>{std::get<1>(*heterogeneous_fan_out_)->as_type<int>(),
-                                        std::get<1>(*heterogeneous_fan_out_)->size_}),
-            options_.return_hops_,
-            options_.with_replacement_,
-            options_.prior_sources_behavior_,
-            options_.dedupe_sources_,
-            do_expensive_check_);
-      } else {
 
-        std::tie(src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets) =
-          cugraph::homogeneous_neighbor_sample(
-            handle_,
-            rng_state_->rng_state_,
-            graph_view,
-            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-            (start_vertex_labels_ != nullptr)
-              ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                    start_vertex_labels->size())
-              : std::nullopt,
-            (label_list_ != nullptr)
-              ? std::make_optional(std::make_tuple(
-                  raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                  label_list_->size_},
-                  raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                  label_to_comm_rank_->size_}))
-              : std::nullopt,
-            raft::host_span<const int>(
-                                      fan_out_->as_type<const int>(), fan_out_->size_),
-            options_.return_hops_,
-            options_.with_replacement_,
-            options_.prior_sources_behavior_,
-            options_.dedupe_sources_,
-            do_expensive_check_);
+      // if 'start_vertex_offsets' is not NULL, leverage the new API
+      if (is_deprecated_) {
+        if (is_biased_) {
+          // Call biased neighbor sample
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
+            cugraph::biased_neighbor_sample(
+              handle_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              
+              (start_vertex_labels_ != nullptr)
+                ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                      start_vertex_labels->size())
+                : std::nullopt,
+              (label_list_ != nullptr)
+                ? std::make_optional(std::make_tuple(
+                    raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                    label_list_->size_},
+                    raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                    label_to_comm_rank_->size_}))
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              rng_state_->rng_state_,
+              options_.return_hops_,
+              options_.with_replacement_,
+              options_.prior_sources_behavior_,
+              options_.dedupe_sources_,
+              do_expensive_check_);
+          
+        } else {
+          // Call uniform neighbor sample
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
+            cugraph::uniform_neighbor_sample(
+              handle_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              
+              (start_vertex_labels_ != nullptr)
+                ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                      start_vertex_labels->size())
+                : std::nullopt,
+              (label_list_ != nullptr)
+                ? std::make_optional(std::make_tuple(
+                    raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                    label_list_->size_},
+                    raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                    label_to_comm_rank_->size_}))
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              rng_state_->rng_state_,
+              options_.return_hops_,
+              options_.with_replacement_,
+              options_.prior_sources_behavior_,
+              options_.dedupe_sources_,
+              do_expensive_check_);
+        }
+
+      } else { // leverage the new API
+
+        rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
 
+        cugraph::detail::sequence_fill(handle_.get_stream(),
+                              vertex_type_offsets.begin(),
+                              vertex_type_offsets.size(),
+                              vertex_t{0});
+
+        if (true) {
+          // call heterogeneous neighbor sample
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::heterogeneous_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              num_edge_types_,
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+        } else {
+          // Call homogeneous neighbor sample
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::homogeneous_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+        }
       }
 
       std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
@@ -357,7 +448,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
-                edge_label ? edge_label->size() : size_t{1},
+                edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
                 hop ? fan_out_->size_ : size_t{1},
                 src_is_major,
                 do_expensive_check_);
@@ -399,7 +490,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
-                edge_label ? edge_label->size() : size_t{1},
+                edge_label ? edge_label->size() : size_t{1},  // FIXME: update edge_label
                 hop ? fan_out_->size_ : size_t{1},
                 src_is_major,
                 options_.compress_per_hop_,
@@ -463,7 +554,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               raft::device_span<vertex_t const>{
                     vertex_type_offsets.data(), vertex_type_offsets.size()},
               
-              edge_label ? edge_label->size() : size_t{1},
+              edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
               hop ? fan_out_->size_ : size_t{1},
               size_t{1},
               size_t{1},
@@ -501,7 +592,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                                            ? std::make_optional(raft::device_span<size_t const>{
                                                offsets->data(), offsets->size()})
                                            : std::nullopt,
-                                         edge_label ? edge_label->size() : size_t{1},
+                                         edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
                                          hop ? fan_out_->size_ : size_t{1},
                                          src_is_major,
                                          do_expensive_check_);
@@ -535,7 +626,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           ? new cugraph::c_api::cugraph_type_erased_device_array_t(*label_hop_offsets, SIZE_T)
           : nullptr,
         (edge_label)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32) // FIXME: update edge_label
           : nullptr,
         (renumber_map) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
                            renumber_map.value(), graph_->vertex_type_)
@@ -1094,13 +1185,15 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                     nullptr,
                                     start_vertices,
                                     start_vertex_labels,
+                                    nullptr, //start_vertex_offsets
                                     label_list,
                                     label_to_comm_rank,
                                     label_offsets,
                                     fan_out,
-                                    nullptr,
+                                    1,
                                     std::move(options_cpp),
                                     is_biased,
+                                    true, // is_deprecated
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1121,6 +1214,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
+
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
   CAPI_EXPECTS(
@@ -1176,13 +1270,15 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_labels,
+                                    nullptr, //start_vertex_offsets
                                     label_list,
                                     label_to_comm_rank,
                                     label_offsets,
                                     fan_out,
-                                    nullptr,
+                                    1,
                                     std::move(options_cpp),
                                     is_biased,
+                                    true, // is_deprecated
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1193,11 +1289,12 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
-  const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
@@ -1219,18 +1316,18 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "must specify label_offsets if retain_seeds is true",
                *error);
-
-  CAPI_EXPECTS((start_vertex_labels == nullptr) ||
+  
+  CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-                    start_vertex_labels)
-                    ->type_ == INT32),
+                    start_vertex_offsets)
+                    ->type_ == SIZE_T),
                CUGRAPH_INVALID_INPUT,
-               "start_vertex_labels should be of type int",
+               "start_vertex_offsets should be of type size_t",
                *error);
 
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_labels != nullptr),
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_labels is also specified",
+               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
                *error);
 
   CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
@@ -1238,19 +1335,12 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                "cannot specify label_to_comm_rank unless label_list is also specified",
                *error);
     
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                  std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
-                  ->type_ == INT32,
-              CUGRAPH_INVALID_INPUT,
-              "edge type offsets type must be INT32",
-              *error);
-  
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                  std::get<0>(*reinterpret_cast<cugraph::c_api::cugraph_sample_heterogeneous_fan_out_t const*>(heterogeneous_fan_out)))
-                  ->type_ == INT32,
-              CUGRAPH_INVALID_INPUT,
-              "fan_out values type must be INT32",
-              *error);
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out should be of type int",
+    *error);
 
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
@@ -1265,14 +1355,16 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    start_vertex_labels,
+                                    nullptr, //start_vertex_labels
+                                    start_vertex_offsets,
                                     label_list,
                                     label_to_comm_rank,
                                     label_offsets,
-                                    nullptr,
-                                    heterogeneous_fan_out,
+                                    fan_out,
+                                    num_edge_types,
                                     std::move(options_cpp),
                                     is_biased,
+                                    false, // is_deprecated
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1283,7 +1375,7 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_device_array_view_t* label_offsets,
@@ -1309,18 +1401,18 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "must specify label_offsets if retain_seeds is true",
                *error);
-
-  CAPI_EXPECTS((start_vertex_labels == nullptr) ||
+  
+  CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-                    start_vertex_labels)
-                    ->type_ == INT32),
+                    start_vertex_offsets)
+                    ->type_ == SIZE_T),
                CUGRAPH_INVALID_INPUT,
-               "start_vertex_labels should be of type int",
+               "start_vertex_offsets should be of type size_t",
                *error);
 
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_labels != nullptr),
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_labels is also specified",
+               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
                *error);
 
   CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
@@ -1349,14 +1441,16 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    start_vertex_labels,
+                                    nullptr, //start_vertex_labels
+                                    start_vertex_offsets,
                                     label_list,
                                     label_to_comm_rank,
                                     label_offsets,
                                     fan_out,
-                                    nullptr,
+                                    1,
                                     std::move(options_cpp),
                                     is_biased,
+                                    false, // is_deprecated
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 6b27b87c7f5..65be8da0b46 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -62,9 +62,8 @@ neighbor_sample_impl(raft::handle_t const& handle,
                      raft::device_span<vertex_t const> starting_vertices,
                      std::optional<raft::device_span<label_t const>> starting_vertex_labels,
                      std::optional<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,
+                     edge_type_t num_edge_types,
                      bool return_hops,
                      bool with_replacement,
                      prior_sources_behavior_t prior_sources_behavior,
@@ -77,54 +76,34 @@ neighbor_sample_impl(raft::handle_t const& handle,
   std::optional<graph_view_t<vertex_t, edge_t, false, multi_gpu>> modified_graph_view{std::nullopt};
 
 
-  if (fan_out) {
-    CUGRAPH_EXPECTS((*fan_out).size() > 0,
-                    "Invalid input argument: number of levels must be non-zero.");
-    CUGRAPH_EXPECTS(
-      (*fan_out).size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
-      "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
-                                                                                // to store hops
-  } else {
-
-    CUGRAPH_EXPECTS(
-      std::get<0>(*heterogeneous_fan_out).back() == std::get<1>(*heterogeneous_fan_out).size() &&
-        std::get<1>(*heterogeneous_fan_out).size() != 0,
-      "Invalid input argument: number of levels and size must match and should be non zero.");
-    
-
-    CUGRAPH_EXPECTS(
-      std::get<0>(*heterogeneous_fan_out).size() <=
-          static_cast<size_t>(std::numeric_limits<int32_t>::max()) &&
-        std::get<1>(*heterogeneous_fan_out).size() <=
-          static_cast<size_t>(std::numeric_limits<int32_t>::max()),
-      "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
-                                                                                // to store hops
+  CUGRAPH_EXPECTS(fan_out.size() > 0, "Invalid input argument: number of levels must be non-zero.");
+  CUGRAPH_EXPECTS(
+    fan_out.size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
+    "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
+                                                                              // to store hops
+  edge_masks_vector.reserve(num_edge_types);
+  
+  for (int i = 0; i < num_edge_types; i++) {
 
-    edge_masks_vector.reserve(std::get<0>(*heterogeneous_fan_out).size() - 1);
+    cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
     
-    for (int i = 0; i < std::get<0>(*heterogeneous_fan_out).size() - 1; i++) {
-
-      cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
-      
-      cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
-
-      (*modified_graph_view).attach_edge_mask(edge_mask.view());
-
-      cugraph::transform_e(
-        handle,
-        *modified_graph_view,
-        cugraph::edge_src_dummy_property_t{}.view(),
-        cugraph::edge_dst_dummy_property_t{}.view(),
-        *edge_type_view,
-        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
-          return edge_type == valid_edge_type;
-        },
-        edge_mask.mutable_view(),
-        false);
-      
-      edge_masks_vector.push_back(std::move(edge_mask));
-    }
+    cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
+
+    (*modified_graph_view).attach_edge_mask(edge_mask.view());
+
+    cugraph::transform_e(
+      handle,
+      *modified_graph_view,
+      cugraph::edge_src_dummy_property_t{}.view(),
+      cugraph::edge_dst_dummy_property_t{}.view(),
+      *edge_type_view,
+      [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
+        return edge_type == valid_edge_type;
+      },
+      edge_mask.mutable_view(),
+      false);
     
+    edge_masks_vector.push_back(std::move(edge_mask));
   }
 
   if constexpr (!multi_gpu) {
@@ -164,14 +143,14 @@ neighbor_sample_impl(raft::handle_t const& handle,
     starting_vertex_labels ? std::make_optional(std::vector<rmm::device_uvector<label_t>>{})
                            : std::nullopt;
 
-  level_result_src_vectors.reserve((*fan_out).size());
-  level_result_dst_vectors.reserve((*fan_out).size());
-  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve((*fan_out).size()); }
-  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve((*fan_out).size()); }
+  level_result_src_vectors.reserve((fan_out).size());
+  level_result_dst_vectors.reserve((fan_out).size());
+  if (level_result_weight_vectors) { (*level_result_weight_vectors).reserve((fan_out).size()); }
+  if (level_result_edge_id_vectors) { (*level_result_edge_id_vectors).reserve((fan_out).size()); }
   if (level_result_edge_type_vectors) {
-    (*level_result_edge_type_vectors).reserve((*fan_out).size());
+    (*level_result_edge_type_vectors).reserve((fan_out).size());
   }
-  if (level_result_label_vectors) { (*level_result_label_vectors).reserve((*fan_out).size()); }
+  if (level_result_label_vectors) { (*level_result_label_vectors).reserve((fan_out).size()); }
 
   rmm::device_uvector<vertex_t> frontier_vertices(0, handle.get_stream());
   auto frontier_vertex_labels =
@@ -192,45 +171,16 @@ neighbor_sample_impl(raft::handle_t const& handle,
   }
 
   std::vector<size_t> level_sizes{};
-  int32_t hop{0};
-  int32_t edge_type_id_max{1}; // A value of 1 translate to homogeneous neighbor sample
-  int32_t num_edge_type_per_hop{0};
   
   auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
 
-  if (heterogeneous_fan_out) {
-    num_edge_type_per_hop = std::get<0>(*heterogeneous_fan_out).back() - 1;
-  }
 
-  while(true) {
-    int32_t k_level{0};
-    if (fan_out) {
-      k_level = (*fan_out)[hop];
-      if (hop == (*fan_out).size()) {
-        break;
-      }
-    } else if (heterogeneous_fan_out) {
-        if (num_edge_type_per_hop == 0) {
-          break;
-        }
-        edge_type_id_max = std::get<0>(*heterogeneous_fan_out).back() - 1;
-      
-    }
-  
-    for (int i = 0; i < edge_type_id_max; i++) {
-      if (heterogeneous_fan_out) { 
-        // Can make the code easier to read by setting a mask for both
-        // homogeneous and heterogeneous neighbor sample. For the former,
-        // no edges should be masked
-        cur_graph_view.attach_edge_mask(edge_masks_vector[i].view());
-        auto k_level_size = (std::get<1>(*heterogeneous_fan_out)[i + 1] - std::get<1>(*heterogeneous_fan_out)[i]);
-        if (k_level_size > hop) {
-          k_level = i + hop;
-        } else { // otherwise, k_level = 0
-          --num_edge_type_per_hop ;
-
-        }
-      }
+  // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1
+  auto num_hops = (fan_out.size() % num_edge_types) ? (fan_out.size() / num_edge_types) : (fan_out.size() / num_edge_types + 1);
+
+  for (auto hop = 0; hop < num_hops; hop++){
+    for (auto edge_type_id = 0; edge_type_id < num_edge_types; edge_type_id++) {
+      auto k_level = fan_out[(hop*num_edge_types) + edge_type_id];
       rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
       rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());
       std::optional<rmm::device_uvector<weight_t>> weights{std::nullopt};
@@ -247,8 +197,8 @@ neighbor_sample_impl(raft::handle_t const& handle,
                       edge_type_view,
                       edge_bias_view,
                       rng_state,
-                      this_frontier_vertices,
-                      this_frontier_vertex_labels,
+                      starting_vertices,
+                      starting_vertex_labels,
                       static_cast<size_t>(k_level),
                       with_replacement);
       } else {
@@ -258,8 +208,8 @@ neighbor_sample_impl(raft::handle_t const& handle,
                                   edge_weight_view,
                                   edge_id_view,
                                   edge_type_view,
-                                  this_frontier_vertices,
-                                  this_frontier_vertex_labels);
+                                  starting_vertices,
+                                  starting_vertex_labels);
       }
 
       level_sizes.push_back(srcs.size());
@@ -272,39 +222,35 @@ neighbor_sample_impl(raft::handle_t const& handle,
       if (labels) { (*level_result_label_vectors).push_back(std::move(*labels)); }
     }
 
-    ++hop;
-    
-    if (hop < (*fan_out).size()) {
-      // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span
-      //  rather than making a copy.
-      auto vertex_partition_range_lasts = cur_graph_view.vertex_partition_range_lasts();
-      std::tie(frontier_vertices, frontier_vertex_labels, vertex_used_as_source) =
-        prepare_next_frontier(
-          handle,
-          starting_vertices,
-          starting_vertex_labels,
-          raft::device_span<vertex_t const>{level_result_dst_vectors.back().data(),
-                                            level_result_dst_vectors.back().size()},
-          frontier_vertex_labels ? std::make_optional(raft::device_span<label_t const>(
-                                     level_result_label_vectors->back().data(),
-                                     level_result_label_vectors->back().size()))
-                                 : std::nullopt,
-          std::move(vertex_used_as_source),
-          cur_graph_view.local_vertex_partition_view(),
-          vertex_partition_range_lasts,
-          prior_sources_behavior,
-          dedupe_sources,
-          do_expensive_check);
-
-      starting_vertices =
-        raft::device_span<vertex_t const>(frontier_vertices.data(), frontier_vertices.size());
-
-      if (frontier_vertex_labels) {
-        starting_vertex_labels = raft::device_span<label_t const>(frontier_vertex_labels->data(),
-                                                                  frontier_vertex_labels->size());
-      }
+    // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span
+    //  rather than making a copy.
+    auto vertex_partition_range_lasts = cur_graph_view.vertex_partition_range_lasts();
+    std::tie(frontier_vertices, frontier_vertex_labels, vertex_used_as_source) =
+      prepare_next_frontier(
+        handle,
+        starting_vertices,
+        starting_vertex_labels,
+        raft::device_span<vertex_t const>{level_result_dst_vectors.back().data(),
+                                          level_result_dst_vectors.back().size()},
+        frontier_vertex_labels ? std::make_optional(raft::device_span<label_t const>(
+                                    level_result_label_vectors->back().data(),
+                                    level_result_label_vectors->back().size()))
+                                : std::nullopt,
+        std::move(vertex_used_as_source),
+        cur_graph_view.local_vertex_partition_view(),
+        vertex_partition_range_lasts,
+        prior_sources_behavior,
+        dedupe_sources,
+        do_expensive_check);
+
+    starting_vertices =
+      raft::device_span<vertex_t const>(frontier_vertices.data(), frontier_vertices.size());
+
+    if (frontier_vertex_labels) {
+      starting_vertex_labels = raft::device_span<label_t const>(frontier_vertex_labels->data(),
+                                                                frontier_vertex_labels->size());
     }
-  } // While loop end here: FIXME: remove this comment
+  }
 
   auto result_size = std::reduce(level_sizes.begin(), level_sizes.end());
   size_t output_offset{};
@@ -386,7 +332,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
   if (return_hops) {
     result_hops   = rmm::device_uvector<int32_t>(result_size, handle.get_stream());
     output_offset = 0;
-    for (size_t i = 0; i < (*fan_out).size(); ++i) {
+    for (size_t i = 0; i < (fan_out).size(); ++i) {
       scalar_fill(
         handle, result_hops->data() + output_offset, level_sizes[i], static_cast<int32_t>(i));
       output_offset += level_sizes[i];
@@ -476,8 +422,8 @@ uniform_neighbor_sample(
     label_to_output_comm_rank
       ? std::make_optional(raft::device_span<int32_t const>{label_map.data(), label_map.size()})
       : std::nullopt,
-    std::make_optional(fan_out),
-    std::nullopt,
+    fan_out,
+    edge_type_t{1},
     return_hops,
     with_replacement,
     prior_sources_behavior,
@@ -540,7 +486,7 @@ biased_neighbor_sample(
       ? std::make_optional(raft::device_span<int32_t const>{label_map.data(), label_map.size()})
       : std::nullopt,
     fan_out,
-    std::nullopt,
+    edge_type_t{1},
     return_hops,
     with_replacement,
     prior_sources_behavior,
@@ -573,8 +519,8 @@ heterogeneous_neighbor_sample(
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  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,
+  edge_type_t num_edge_types,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 { 
@@ -583,26 +529,37 @@ heterogeneous_neighbor_sample(
   if (starting_vertex_offsets)
     starting_vertex_labels =
       detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
-
-  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    rng_state,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    edge_bias_view,
-    starting_vertices,
-    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
-                                starting_vertex_labels.data(), starting_vertex_labels.size()})
-                            : std::nullopt,
-    label_to_output_comm_rank,
-    fan_out,
-    sampling_flags.return_hops,
-    sampling_flags.with_replacement,
-    sampling_flags.prior_sources_behavior,
-    sampling_flags.dedupe_sources,
-    do_expensive_check);
+  
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
+    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        edge_id_view,
+        edge_type_view,
+        edge_bias_view,
+        starting_vertices,
+        starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
+                                    starting_vertex_labels.data(), starting_vertex_labels.size()})
+                                : std::nullopt,
+        label_to_output_comm_rank,
+        fan_out,
+        num_edge_types,
+        sampling_flags.return_hops,
+        sampling_flags.with_replacement,
+        sampling_flags.prior_sources_behavior,
+        sampling_flags.dedupe_sources,
+        do_expensive_check);
+  
+  return std::make_tuple(std::move(majors),
+                         std::move(minors),
+                         std::move(weights),
+                         std::move(edge_ids),
+                         std::move(edge_types),
+                         std::move(hops),
+                         std::move(offsets));
+  
 }
 
 template <typename vertex_t,
@@ -631,6 +588,7 @@ homogeneous_neighbor_sample(
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
+  // edge_type_t num_edge_types, /*argument not needed for homogeneous neighbor sample*/
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
@@ -641,8 +599,11 @@ homogeneous_neighbor_sample(
     starting_vertex_labels =
       detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
 
-  return detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+  
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
+    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
+    rng_state,
     graph_view,
     edge_weight_view,
     edge_id_view,
@@ -654,13 +615,20 @@ homogeneous_neighbor_sample(
                             : std::nullopt,
     label_to_output_comm_rank,
     fan_out,
-    std::nullopt,
-    return_hops,
-    with_replacement,
-    prior_sources_behavior,
-    dedupe_sources,
-    rng_state,
+    edge_type_t{1},
+    sampling_flags.return_hops,
+    sampling_flags.with_replacement,
+    sampling_flags.prior_sources_behavior,
+    sampling_flags.dedupe_sources,
     do_expensive_check);
+  
+  return std::make_tuple(std::move(majors),
+                         std::move(minors),
+                         std::move(weights),
+                         std::move(edge_ids),
+                         std::move(edge_types),
+                         std::move(hops),
+                         std::move(offsets));
 }
 
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
index 917e8fcc6e1..99b0ea82ea5 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
@@ -133,7 +133,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -144,15 +143,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -161,7 +156,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -172,15 +166,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -189,7 +179,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -200,14 +189,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -216,7 +201,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -227,14 +211,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
index 71661090023..159c2b3c372 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
@@ -134,7 +134,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -145,15 +144,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -162,7 +157,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -173,15 +167,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -190,7 +180,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -201,14 +190,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -217,7 +202,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -228,14 +212,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
index 2a1c0d43ab3..0d69815778f 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
@@ -133,7 +133,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -144,15 +143,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -161,7 +156,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -172,15 +166,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -189,7 +179,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -200,14 +189,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -216,7 +201,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -227,14 +211,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
index c58e1659d7c..25a597b53fd 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
@@ -133,7 +133,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -144,15 +143,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -161,7 +156,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -172,15 +166,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -189,7 +179,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -200,14 +189,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -216,7 +201,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -227,14 +211,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
index 2a70085d490..feb1c3cf052 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
@@ -133,7 +133,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -144,15 +143,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -161,7 +156,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -172,15 +166,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -189,7 +179,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -200,14 +189,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
@@ -216,7 +201,6 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -227,14 +211,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
index 6d045ad092d..334fef6ff8f 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
@@ -133,7 +133,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -144,15 +143,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -161,7 +156,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -172,15 +166,11 @@ heterogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::tuple<raft::host_span<int32_t const>, raft::host_span<int32_t const>>
-    heterogeneous_fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -189,7 +179,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -200,14 +189,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int64_t>,
@@ -216,7 +201,6 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int64_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
 homogeneous_neighbor_sample(
   raft::handle_t const& handle,
@@ -227,14 +211,10 @@ homogeneous_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  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::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
+  sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
 }  // namespace cugraph

From 0358c6eff3fdb79f2a0e08886df0990c8a556743 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 09:43:27 -0700
Subject: [PATCH 039/117] add fixme for deprecated flags

---
 cpp/src/c_api/neighbor_sampling.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 62f55c19472..42961f32ee2 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -1184,7 +1184,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                     graph,
                                     nullptr,
                                     start_vertices,
-                                    start_vertex_labels,
+                                    start_vertex_labels, // Deprecated flag
                                     nullptr, //start_vertex_offsets
                                     label_list,
                                     label_to_comm_rank,
@@ -1269,7 +1269,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    start_vertex_labels,
+                                    start_vertex_labels, // Deprecated flag
                                     nullptr, //start_vertex_offsets
                                     label_list,
                                     label_to_comm_rank,
@@ -1355,7 +1355,7 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    nullptr, //start_vertex_labels
+                                    nullptr, //start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
                                     label_list,
                                     label_to_comm_rank,
@@ -1441,7 +1441,7 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    nullptr, //start_vertex_labels
+                                    nullptr, //start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
                                     label_list,
                                     label_to_comm_rank,

From 799c35dc5ebcdd74e8ed545bcff1c338c767d4e6 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 16:31:58 -0700
Subject: [PATCH 040/117] update CAPI

---
 cpp/include/cugraph_c/sampling_algorithms.h | 44 ++++++----------
 cpp/src/c_api/neighbor_sampling.cpp         | 56 +++++++++------------
 2 files changed, 38 insertions(+), 62 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index b9e7836b550..1b2d15dea12 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -509,20 +509,14 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
  * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_labels  Device array of start vertex labels for the sampling.  The
- * labels associated with each start vertex will be included in the output associated with results
- * that were derived from that start vertex.  We only support label of type INT32. If label is
- * NULL, the return data will not be labeled.
- * @param [in]  label_list Device array of the labels included in @p start_vertex_labels.  If
- * @p label_to_comm_rank is not specified this parameter is ignored.  If specified, label_list
- * must be sorted in ascending order.
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
  * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output.  If not specifed the data is not organized in
- * output.  If specified then the all data from @p label_list[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_labels is also specified
+ * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
+ * If not specifed the data is not organized in output.  If specified then the all data in the range
+ * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_offsets is also specified.
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
- * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
@@ -546,10 +540,8 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
-  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_sample_heterogeneous_fan_out_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -577,20 +569,14 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
  * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
  * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_labels  Device array of start vertex labels for the sampling.  The
- * labels associated with each start vertex will be included in the output associated with results
- * that were derived from that start vertex.  We only support label of type INT32. If label is
- * NULL, the return data will not be labeled.
- * @param [in]  label_list Device array of the labels included in @p start_vertex_labels.  If
- * @p label_to_comm_rank is not specified this parameter is ignored.  If specified, label_list
- * must be sorted in ascending order.
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
  * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output.  If not specifed the data is not organized in
- * output.  If specified then the all data from @p label_list[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_labels is also specified
+ * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
+ * If not specifed the data is not organized in output.  If specified then the all data in the range
+ * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_offsets is also specified.
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
- * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  sampling_options
@@ -611,10 +597,8 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_labels,
-  const cugraph_type_erased_device_array_view_t* label_list,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 42961f32ee2..77414028d7b 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -88,7 +88,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   int num_edge_types_{};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool is_biased_{false};
-  bool is_deprecated_{false};
+  bool is_deprecated_api_{false};
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
@@ -107,7 +107,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     int num_edge_types,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool is_biased,
-    bool is_deprecated,
+    bool is_deprecated_api,
     bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
@@ -137,7 +137,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       num_edge_types_(num_edge_types),
       options_(options),
       is_biased_(is_biased),
-      is_deprecated_(is_deprecated),
+      is_deprecated_api_(is_deprecated_api),
       do_expensive_check_(do_expensive_check)
   {
   }
@@ -250,7 +250,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
 
       // if 'start_vertex_offsets' is not NULL, leverage the new API
-      if (is_deprecated_) {
+      if (is_deprecated_api_) {
         if (is_biased_) {
           // Call biased neighbor sample
           auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
@@ -554,7 +554,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               raft::device_span<vertex_t const>{
                     vertex_type_offsets.data(), vertex_type_offsets.size()},
               
-              edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
+              start_vertex_offsets ? start_vertex_offsets->size() : size_t{1},
               hop ? fan_out_->size_ : size_t{1},
               size_t{1},
               size_t{1},
@@ -580,6 +580,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           CUGRAPH_FAIL("Can only use COO format if not renumbering");
         }
 
+        if (num_edge_types_ != 1) {
+          CUGRAPH_FAIL("Can only use COO format for homogeneous neighborhood sampling");
+        }
+
         std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
           cugraph::sort_sampled_edgelist(handle_,
                                          std::move(src),
@@ -1193,7 +1197,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
                                     1,
                                     std::move(options_cpp),
                                     is_biased,
-                                    true, // is_deprecated
+                                    true, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1278,7 +1282,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                     1,
                                     std::move(options_cpp),
                                     is_biased,
-                                    true, // is_deprecated
+                                    true, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1290,9 +1294,7 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -1312,9 +1314,10 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
       *error);
   }
   
-  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
+  // FIXME: Should we maintain this contition?
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
-               "must specify label_offsets if retain_seeds is true",
+               "must specify start_vertex_offsets if retain_seeds is true",
                *error);
   
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
@@ -1329,11 +1332,6 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless label_list is also specified",
-               *error);
     
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
@@ -1355,16 +1353,16 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    nullptr, //start_vertex_labels /*Deprecated flag*/
+                                    nullptr, // start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
-                                    label_list,
+                                    nullptr, // label_list /*Deprecated flag*/
                                     label_to_comm_rank,
-                                    label_offsets,
+                                    nullptr, // label_offsets /*Deprecated flag*/
                                     fan_out,
                                     num_edge_types,
                                     std::move(options_cpp),
                                     is_biased,
-                                    false, // is_deprecated
+                                    false, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1376,9 +1374,7 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_list,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-  const cugraph_type_erased_device_array_view_t* label_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
@@ -1397,9 +1393,10 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
       *error);
   }
   
-  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (label_offsets != nullptr),
+  // FIXME: Should we maintain this contition?
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
-               "must specify label_offsets if retain_seeds is true",
+               "must specify start_vertex_offsets if retain_seeds is true",
                *error);
   
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
@@ -1414,11 +1411,6 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (label_list != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless label_list is also specified",
-               *error);
   
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
                 fan_out)
@@ -1443,14 +1435,14 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                                     start_vertices,
                                     nullptr, //start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
-                                    label_list,
+                                    nullptr,//label_list /*Deprecated flag*/
                                     label_to_comm_rank,
-                                    label_offsets,
+                                    nullptr, //label_offsets /*Deprecated flag*/
                                     fan_out,
                                     1,
                                     std::move(options_cpp),
                                     is_biased,
-                                    false, // is_deprecated
+                                    false, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }

From ab8aa72ac167aef50a949c0ff49367d36e35b80d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 17:08:35 -0700
Subject: [PATCH 041/117] undo changes to k-truss

---
 cpp/src/c_api/neighbor_sampling.cpp | 25 ++++++++++++++-----------
 cpp/src/community/k_truss_impl.cuh  | 19 ++++---------------
 2 files changed, 18 insertions(+), 26 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 77414028d7b..d708b6fbcc2 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -416,7 +416,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                           (options_.compression_type_ == cugraph_compression_type_t::COO);
 
       if (options_.renumber_results_) {
-        if (fan_out_ != nullptr) {
+        if (num_edge_types_ == 1) { // homogeneous renumbering
           if (options_.compression_type_ == cugraph_compression_type_t::COO) {
             // COO
 
@@ -442,13 +442,14 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                options_.retain_seeds_ ? (is_deprecated_api_? std::make_optional(raft::device_span<size_t const>{
+                                          label_offsets_->as_type<size_t>(), label_offsets_->size_}) : std::make_optional(raft::device_span<size_t const>{
+                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_}))
                                       : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
-                edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
+                edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label ?
                 hop ? fan_out_->size_ : size_t{1},
                 src_is_major,
                 do_expensive_check_);
@@ -463,6 +464,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
             rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
             rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+            // FIXME: Update this function to handle the new API with 'starting_vertex_offsets'
+            // and indices.
             std::tie(majors,
                     output_major_offsets,
                     minors,
@@ -484,8 +487,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                          label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                options_.retain_seeds_ ? (is_deprecated_api_? std::make_optional(raft::device_span<size_t const>{
+                                          label_offsets_->as_type<size_t>(), label_offsets_->size_}) : std::make_optional(raft::device_span<size_t const>{
+                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_}))
                                       : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
@@ -505,7 +509,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           hop.reset();
           offsets.reset();
         
-        } else { // heterogeneous
+        } else { // heterogeneous renumbering
 
           rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
 
@@ -523,6 +527,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           // extract the edge_type from label_type_hop_offsets
           std::optional<rmm::device_uvector<size_t>> label_type_hop_offsets{std::nullopt};
 
+          // FIXME: Update this function to handle the new API with 'starting_vertex_offsets'
+          // and indices.
           std::tie(output_majors,
               minors,
               wgt,
@@ -546,7 +552,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                     start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                 : std::nullopt,
               options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                        label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                        start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
                                     : std::nullopt,
               offsets ? std::make_optional(
                           raft::device_span<size_t const>{offsets->data(), offsets->size()})
@@ -580,9 +586,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           CUGRAPH_FAIL("Can only use COO format if not renumbering");
         }
 
-        if (num_edge_types_ != 1) {
-          CUGRAPH_FAIL("Can only use COO format for homogeneous neighborhood sampling");
-        }
 
         std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
           cugraph::sort_sampled_edgelist(handle_,
diff --git a/cpp/src/community/k_truss_impl.cuh b/cpp/src/community/k_truss_impl.cuh
index 4facf093535..e052a892917 100644
--- a/cpp/src/community/k_truss_impl.cuh
+++ b/cpp/src/community/k_truss_impl.cuh
@@ -148,13 +148,13 @@ k_truss(raft::handle_t const& handle,
         std::nullopt,
         std::nullopt,
         cugraph::graph_properties_t{true, graph_view.is_multigraph()},
-        false);
+        true);
 
     modified_graph_view = (*modified_graph).view();
   }
 
   // 2. Find (k-1)-core and exclude edges that do not belong to (k-1)-core
-  /*
+
   {
     auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
 
@@ -211,7 +211,6 @@ k_truss(raft::handle_t const& handle,
 
     renumber_map = std::move(tmp_renumber_map);
   }
-  */
 
   // 3. Keep only the edges from a low-degree vertex to a high-degree vertex.
 
@@ -278,7 +277,7 @@ k_truss(raft::handle_t const& handle,
         std::nullopt,
         std::nullopt,
         cugraph::graph_properties_t{false /* now asymmetric */, cur_graph_view.is_multigraph()},
-        false);
+        true);
 
     modified_graph_view = (*modified_graph).view();
     if (renumber_map) {  // collapse renumber_map
@@ -341,12 +340,8 @@ k_truss(raft::handle_t const& handle,
         edge_weight_view ? std::make_optional(*edge_weight_view) : std::nullopt,
         std::optional<edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        /*
         std::make_optional(
-          raft::device_span<vertex_t const>((*renumber_map).data(), (*renumber_map).size()))
-        */
-        std::optional<raft::device_span<vertex_t const>>(std::nullopt)
-        );
+          raft::device_span<vertex_t const>((*renumber_map).data(), (*renumber_map).size())));
 
     std::tie(edgelist_srcs, edgelist_dsts, edgelist_wgts) =
       symmetrize_edgelist<vertex_t, weight_t, false, multi_gpu>(handle,
@@ -354,12 +349,6 @@ k_truss(raft::handle_t const& handle,
                                                                 std::move(edgelist_dsts),
                                                                 std::move(edgelist_wgts),
                                                                 false);
-    
-    raft::print_device_vector("edgelist_srcs", edgelist_srcs.data(), edgelist_srcs.size(), std::cout);
-    raft::print_device_vector("edgelist_dsts", edgelist_dsts.data(), edgelist_dsts.size(), std::cout);
-
-    printf("\nK-TRUSS Successfully completed and the subgraph size = %d\n", edgelist_srcs.size());
-    //raft::print_device_vector("edgelist_wgts", count_3_.data(), count_3_.size(), std::cout); 
 
     return std::make_tuple(
       std::move(edgelist_srcs), std::move(edgelist_dsts), std::move(edgelist_wgts));

From 7d8b5add63d27ac4dfd77ff085bddff3a7d0e62d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 17:10:12 -0700
Subject: [PATCH 042/117] undo changes to tests

---
 cpp/tests/community/k_truss_test.cpp | 54 ++++++++++++++++++++++++----
 1 file changed, 48 insertions(+), 6 deletions(-)

diff --git a/cpp/tests/community/k_truss_test.cpp b/cpp/tests/community/k_truss_test.cpp
index af34958f0de..424d52f2067 100644
--- a/cpp/tests/community/k_truss_test.cpp
+++ b/cpp/tests/community/k_truss_test.cpp
@@ -262,6 +262,7 @@ class Tests_KTruss : public ::testing::TestWithParam<std::tuple<KTruss_Usecase,
 };
 
 using Tests_KTruss_File = Tests_KTruss<cugraph::test::File_Usecase>;
+using Tests_KTruss_Rmat = Tests_KTruss<cugraph::test::Rmat_Usecase>;
 
 TEST_P(Tests_KTruss_File, CheckInt32Int32Float)
 {
@@ -269,16 +270,57 @@ TEST_P(Tests_KTruss_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
+TEST_P(Tests_KTruss_File, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_KTruss_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_KTruss_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
 INSTANTIATE_TEST_SUITE_P(
   simple_test,
   Tests_KTruss_File,
   ::testing::Combine(
     // enable correctness checks
-    ::testing::Values(KTruss_Usecase{5, false, false}),
-    ::testing::Values(
-      cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/datasets/test_datasets_.mtx")
-      //cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/cugraph/datasets/dolphins.mtx"),
-      //cugraph::test::File_Usecase("/raid/jnke/optimize_ktruss/cugraph/datasets/netscience.mtx")
-      )));
+    ::testing::Values(KTruss_Usecase{5, true, false},
+                      KTruss_Usecase{4, true, false},
+                      KTruss_Usecase{9, true, true},
+                      KTruss_Usecase{7, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/netscience.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/dolphins.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(rmat_small_test,
+                         Tests_KTruss_Rmat,
+                         // enable correctness checks
+                         ::testing::Combine(::testing::Values(KTruss_Usecase{5, false, true},
+                                                              KTruss_Usecase{4, false, true},
+                                                              KTruss_Usecase{9, true, true},
+                                                              KTruss_Usecase{7, true, true}),
+                                            ::testing::Values(cugraph::test::Rmat_Usecase(
+                                              10, 16, 0.57, 0.19, 0.19, 0, true, false))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_KTruss_Rmat,
+  // disable correctness checks for large graphs
+  // FIXME: High memory footprint. Perform nbr_intersection in chunks.
+  ::testing::Combine(
+    ::testing::Values(KTruss_Usecase{12, false, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(14, 16, 0.57, 0.19, 0.19, 0, true, false))));
 
 CUGRAPH_TEST_PROGRAM_MAIN()

From f2190bae3a2cf93f6e92bac495479126383964cc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 20 Sep 2024 18:32:31 -0700
Subject: [PATCH 043/117] clean up code

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 65be8da0b46..cd62bbd06c5 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -73,7 +73,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
   static_assert(std::is_floating_point_v<bias_t>);
 
   std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
-  std::optional<graph_view_t<vertex_t, edge_t, false, multi_gpu>> modified_graph_view{std::nullopt};
+  graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
 
 
   CUGRAPH_EXPECTS(fan_out.size() > 0, "Invalid input argument: number of levels must be non-zero.");
@@ -87,13 +87,13 @@ neighbor_sample_impl(raft::handle_t const& handle,
 
     cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
     
-    cugraph::fill_edge_property(handle, graph_view, edge_mask.mutable_view(), bool{true});
+    cugraph::fill_edge_property(handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
 
-    (*modified_graph_view).attach_edge_mask(edge_mask.view());
+    modified_graph_view.attach_edge_mask(edge_mask.view());
 
     cugraph::transform_e(
       handle,
-      *modified_graph_view,
+      modified_graph_view,
       cugraph::edge_src_dummy_property_t{}.view(),
       cugraph::edge_dst_dummy_property_t{}.view(),
       *edge_type_view,
@@ -172,7 +172,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
 
   std::vector<size_t> level_sizes{};
   
-  auto cur_graph_view = modified_graph_view ? *modified_graph_view : graph_view;
+  auto cur_graph_view = modified_graph_view;
 
 
   // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1

From 1e96dcff96557aa1ffbd3387964c651a27b31b8c Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 23 Sep 2024 09:34:33 -0700
Subject: [PATCH 044/117] update docs

---
 cpp/include/cugraph/sampling_functions.hpp | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index c59791d4a62..1f82ef67127 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -42,9 +42,6 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
 
 /**
  * @brief Uniform Neighborhood Sampling.
- * 
- * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
- * 'is_biased' set to false instead
  *
  * @deprecated Replaced with homogeneous_uniform_neighbor_sample
  *
@@ -144,9 +141,6 @@ uniform_neighbor_sample(
 
 /**
  * @brief Biased Neighborhood Sampling.
- * 
- * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
- * 'is_biased' set to true instead
  *
  * @deprecated Replaced with homogeneous_biased_neighbor_sample
  *

From 36c25ad17713d90ba97c829852c5fe2726d44bcf Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 23 Sep 2024 11:43:48 -0700
Subject: [PATCH 045/117] fix typo

---
 cpp/src/c_api/neighbor_sampling.cpp           | 1008 +++++++++++++----
 cpp/src/sampling/neighbor_sampling_impl.hpp   |   60 +-
 .../_cugraph_c/sampling_algorithms.pxd        |    6 +-
 3 files changed, 826 insertions(+), 248 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index d708b6fbcc2..a084e0fb5d7 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -22,6 +22,7 @@
 #include "c_api/random.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
+#include "c_api/graph_helper.hpp"
 
 #include <cugraph_c/algorithms.h>
 #include <cugraph_c/sampling_algorithms.h>
@@ -30,6 +31,7 @@
 #include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/sampling_functions.hpp>
+#include "sampling/detail/sampling_utils.hpp"
 
 #include <raft/core/handle.hpp>
 
@@ -73,22 +75,710 @@ struct cugraph_sample_result_t {
 
 namespace {
 
+// Deprecated functor
+struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
+  raft::handle_t const& handle_;
+  cugraph::c_api::cugraph_graph_t* graph_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
+  cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
+  cugraph::c_api::cugraph_sampling_options_t options_{};
+  bool do_expensive_check_{false};
+  cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
+
+  uniform_neighbor_sampling_functor(
+    cugraph_resource_handle_t const* handle,
+    cugraph_graph_t* graph,
+    cugraph_type_erased_device_array_view_t const* start_vertices,
+    cugraph_type_erased_device_array_view_t const* start_vertex_labels,
+    cugraph_type_erased_device_array_view_t const* label_list,
+    cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
+    cugraph_type_erased_device_array_view_t const* label_offsets,
+    cugraph_type_erased_host_array_view_t const* fan_out,
+    cugraph_rng_state_t* rng_state,
+    cugraph::c_api::cugraph_sampling_options_t options,
+    bool do_expensive_check)
+    : abstract_functor(),
+      handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      start_vertices_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertices)),
+      start_vertex_labels_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertex_labels)),
+      label_list_(reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+        label_list)),
+      label_to_comm_rank_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          label_to_comm_rank)),
+      label_offsets_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          label_offsets)),
+      fan_out_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
+      options_(options),
+      do_expensive_check_(do_expensive_check)
+  {
+  }
+
+  template <typename vertex_t,
+            typename edge_t,
+            typename weight_t,
+            typename edge_type_t,
+            bool store_transposed,
+            bool multi_gpu>
+  void operator()()
+  {
+    using label_t = int32_t;
+
+    // FIXME: Think about how to handle SG vice MG
+    if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
+      unsupported();
+    } else {
+      // uniform_nbr_sample expects store_transposed == false
+      if constexpr (store_transposed) {
+        error_code_ = cugraph::c_api::
+          transpose_storage<vertex_t, edge_t, weight_t, store_transposed, multi_gpu>(
+            handle_, graph_, error_.get());
+        if (error_code_ != CUGRAPH_SUCCESS) return;
+      }
+
+      auto graph =
+        reinterpret_cast<cugraph::graph_t<vertex_t, edge_t, false, multi_gpu>*>(graph_->graph_);
+
+      auto graph_view = graph->view();
+
+      auto edge_weights = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 weight_t>*>(graph_->edge_weights_);
+
+      auto edge_ids = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 edge_t>*>(graph_->edge_ids_);
+
+      auto edge_types = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 edge_type_t>*>(graph_->edge_types_);
+
+      auto number_map = reinterpret_cast<rmm::device_uvector<vertex_t>*>(graph_->number_map_);
+
+      rmm::device_uvector<vertex_t> start_vertices(start_vertices_->size_, handle_.get_stream());
+      raft::copy(start_vertices.data(),
+                 start_vertices_->as_type<vertex_t>(),
+                 start_vertices.size(),
+                 handle_.get_stream());
+
+      std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
+
+      if (start_vertex_labels_ != nullptr) {
+        start_vertex_labels =
+          rmm::device_uvector<label_t>{start_vertex_labels_->size_, handle_.get_stream()};
+        raft::copy(start_vertex_labels->data(),
+                   start_vertex_labels_->as_type<label_t>(),
+                   start_vertex_labels_->size_,
+                   handle_.get_stream());
+      }
+
+      if constexpr (multi_gpu) {
+        if (start_vertex_labels) {
+          std::tie(start_vertices, *start_vertex_labels) =
+            cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices), std::move(*start_vertex_labels));
+        } else {
+          start_vertices =
+            cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices));
+        }
+      }
+
+      //
+      // Need to renumber start_vertices
+      //
+      cugraph::renumber_local_ext_vertices<vertex_t, multi_gpu>(
+        handle_,
+        start_vertices.data(),
+        start_vertices.size(),
+        number_map->data(),
+        graph_view.local_vertex_partition_range_first(),
+        graph_view.local_vertex_partition_range_last(),
+        do_expensive_check_);
+
+      auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
+        cugraph::uniform_neighbor_sample(
+          handle_,
+          graph_view,
+          (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+          (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+          (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+          raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+          (start_vertex_labels_ != nullptr)
+            ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                   start_vertex_labels->size())
+            : std::nullopt,
+          (label_list_ != nullptr)
+            ? std::make_optional(std::make_tuple(
+                raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                 label_list_->size_},
+                raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                 label_to_comm_rank_->size_}))
+            : std::nullopt,
+          raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
+          rng_state_->rng_state_,
+          options_.return_hops_,
+          options_.with_replacement_,
+          options_.prior_sources_behavior_,
+          options_.dedupe_sources_,
+          do_expensive_check_);
+
+      std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
+
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
+                                                            src.data(),
+                                                            src.size(),
+                                                            number_map->data(),
+                                                            vertex_partition_lasts,
+                                                            do_expensive_check_);
+
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
+                                                            dst.data(),
+                                                            dst.size(),
+                                                            number_map->data(),
+                                                            vertex_partition_lasts,
+                                                            do_expensive_check_);
+
+      std::optional<rmm::device_uvector<vertex_t>> majors{std::nullopt};
+      rmm::device_uvector<vertex_t> minors(0, handle_.get_stream());
+      std::optional<rmm::device_uvector<size_t>> major_offsets{std::nullopt};
+
+      std::optional<rmm::device_uvector<size_t>> label_hop_offsets{std::nullopt};
+
+      std::optional<rmm::device_uvector<vertex_t>> renumber_map{std::nullopt};
+      std::optional<rmm::device_uvector<size_t>> renumber_map_offsets{std::nullopt};
+
+      bool src_is_major = (options_.compression_type_ == cugraph_compression_type_t::CSR) ||
+                          (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+                          (options_.compression_type_ == cugraph_compression_type_t::COO);
+
+      if (options_.renumber_results_) {
+        if (options_.compression_type_ == cugraph_compression_type_t::COO) {
+          // COO
+
+          rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
+          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+          std::tie(output_majors,
+                   minors,
+                   wgt,
+                   edge_id,
+                   edge_type,
+                   label_hop_offsets,
+                   output_renumber_map,
+                   renumber_map_offsets) =
+            cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
+              handle_,
+              std::move(src),
+              std::move(dst),
+              std::move(wgt),
+              std::move(edge_id),
+              std::move(edge_type),
+              std::move(hop),
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
+              offsets ? std::make_optional(
+                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                      : std::nullopt,
+              edge_label ? edge_label->size() : size_t{1},
+              hop ? fan_out_->size_ : size_t{1},
+              src_is_major,
+              do_expensive_check_);
+
+          majors.emplace(std::move(output_majors));
+          renumber_map.emplace(std::move(output_renumber_map));
+        } else {
+          // (D)CSC, (D)CSR
+
+          bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+                                 (options_.compression_type_ == cugraph_compression_type_t::DCSC);
+
+          rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
+          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+          std::tie(majors,
+                   output_major_offsets,
+                   minors,
+                   wgt,
+                   edge_id,
+                   edge_type,
+                   label_hop_offsets,
+                   output_renumber_map,
+                   renumber_map_offsets) =
+            cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
+              handle_,
+              std::move(src),
+              std::move(dst),
+              std::move(wgt),
+              std::move(edge_id),
+              std::move(edge_type),
+              std::move(hop),
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
+              offsets ? std::make_optional(
+                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                      : std::nullopt,
+              edge_label ? edge_label->size() : size_t{1},
+              hop ? fan_out_->size_ : size_t{1},
+              src_is_major,
+              options_.compress_per_hop_,
+              doubly_compress,
+              do_expensive_check_);
+
+          renumber_map.emplace(std::move(output_renumber_map));
+          major_offsets.emplace(std::move(output_major_offsets));
+        }
+
+        // These are now represented by label_hop_offsets
+        hop.reset();
+        offsets.reset();
+      } else {
+        if (options_.compression_type_ != cugraph_compression_type_t::COO) {
+          CUGRAPH_FAIL("Can only use COO format if not renumbering");
+        }
+
+        std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
+          cugraph::sort_sampled_edgelist(handle_,
+                                         std::move(src),
+                                         std::move(dst),
+                                         std::move(wgt),
+                                         std::move(edge_id),
+                                         std::move(edge_type),
+                                         std::move(hop),
+                                         offsets
+                                           ? std::make_optional(raft::device_span<size_t const>{
+                                               offsets->data(), offsets->size()})
+                                           : std::nullopt,
+                                         edge_label ? edge_label->size() : size_t{1},
+                                         hop ? fan_out_->size_ : size_t{1},
+                                         src_is_major,
+                                         do_expensive_check_);
+
+        majors.emplace(std::move(src));
+        minors = std::move(dst);
+
+        hop.reset();
+        offsets.reset();
+      }
+
+      result_ = new cugraph::c_api::cugraph_sample_result_t{
+        (major_offsets)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*major_offsets, SIZE_T)
+          : nullptr,
+        (majors)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*majors, graph_->vertex_type_)
+          : nullptr,
+        new cugraph::c_api::cugraph_type_erased_device_array_t(minors, graph_->vertex_type_),
+        (edge_id)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*edge_id, graph_->edge_type_)
+          : nullptr,
+        (edge_type) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                        *edge_type, graph_->edge_type_id_type_)
+                    : nullptr,
+        (wgt) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*wgt, graph_->weight_type_)
+              : nullptr,
+        (hop) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*hop, INT32)
+              : nullptr,  // FIXME get rid of this
+        (label_hop_offsets)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*label_hop_offsets, SIZE_T)
+          : nullptr,
+        (edge_label)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32)
+          : nullptr,
+        (renumber_map) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                           renumber_map.value(), graph_->vertex_type_)
+                       : nullptr,
+        (renumber_map_offsets) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                                   renumber_map_offsets.value(), SIZE_T)
+                               : nullptr};
+    }
+  }
+};
+
+// Deprecated functor
+struct biased_neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
+  raft::handle_t const& handle_;
+  cugraph::c_api::cugraph_graph_t* graph_{nullptr};
+  cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
+  cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
+  cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
+  cugraph::c_api::cugraph_sampling_options_t options_{};
+  bool do_expensive_check_{false};
+  cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
+
+  biased_neighbor_sampling_functor(
+    cugraph_resource_handle_t const* handle,
+    cugraph_graph_t* graph,
+    cugraph_edge_property_view_t const* edge_biases,
+    cugraph_type_erased_device_array_view_t const* start_vertices,
+    cugraph_type_erased_device_array_view_t const* start_vertex_labels,
+    cugraph_type_erased_device_array_view_t const* label_list,
+    cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
+    cugraph_type_erased_device_array_view_t const* label_offsets,
+    cugraph_type_erased_host_array_view_t const* fan_out,
+    cugraph_rng_state_t* rng_state,
+    cugraph::c_api::cugraph_sampling_options_t options,
+    bool do_expensive_check)
+    : abstract_functor(),
+      handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
+      graph_(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)),
+      edge_biases_(
+        reinterpret_cast<cugraph::c_api::cugraph_edge_property_view_t const*>(edge_biases)),
+      start_vertices_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertices)),
+      start_vertex_labels_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          start_vertex_labels)),
+      label_list_(reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+        label_list)),
+      label_to_comm_rank_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          label_to_comm_rank)),
+      label_offsets_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+          label_offsets)),
+      fan_out_(
+        reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
+      rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
+      options_(options),
+      do_expensive_check_(do_expensive_check)
+  {
+  }
+
+  template <typename vertex_t,
+            typename edge_t,
+            typename weight_t,
+            typename edge_type_t,
+            bool store_transposed,
+            bool multi_gpu>
+  void operator()()
+  {
+    using label_t = int32_t;
+
+    // FIXME: Think about how to handle SG vice MG
+    if constexpr (!cugraph::is_candidate<vertex_t, edge_t, weight_t>::value) {
+      unsupported();
+    } else {
+      // uniform_nbr_sample expects store_transposed == false
+      if constexpr (store_transposed) {
+        error_code_ = cugraph::c_api::
+          transpose_storage<vertex_t, edge_t, weight_t, store_transposed, multi_gpu>(
+            handle_, graph_, error_.get());
+        if (error_code_ != CUGRAPH_SUCCESS) return;
+      }
+
+      auto graph =
+        reinterpret_cast<cugraph::graph_t<vertex_t, edge_t, false, multi_gpu>*>(graph_->graph_);
+
+      auto graph_view = graph->view();
+
+      auto edge_weights = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 weight_t>*>(graph_->edge_weights_);
+
+      auto edge_ids = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 edge_t>*>(graph_->edge_ids_);
+
+      auto edge_types = reinterpret_cast<
+        cugraph::edge_property_t<cugraph::graph_view_t<vertex_t, edge_t, true, multi_gpu>,
+                                 edge_type_t>*>(graph_->edge_types_);
+
+      auto number_map = reinterpret_cast<rmm::device_uvector<vertex_t>*>(graph_->number_map_);
+
+      auto edge_biases =
+        edge_biases_ ? reinterpret_cast<cugraph::edge_property_view_t<edge_t, weight_t const*>*>(
+                         edge_biases_->edge_property_)
+                     : nullptr;
+
+      rmm::device_uvector<vertex_t> start_vertices(start_vertices_->size_, handle_.get_stream());
+      raft::copy(start_vertices.data(),
+                 start_vertices_->as_type<vertex_t>(),
+                 start_vertices.size(),
+                 handle_.get_stream());
+
+      std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
+
+      if (start_vertex_labels_ != nullptr) {
+        start_vertex_labels =
+          rmm::device_uvector<label_t>{start_vertex_labels_->size_, handle_.get_stream()};
+        raft::copy(start_vertex_labels->data(),
+                   start_vertex_labels_->as_type<label_t>(),
+                   start_vertex_labels_->size_,
+                   handle_.get_stream());
+      }
+
+      if constexpr (multi_gpu) {
+        if (start_vertex_labels) {
+          std::tie(start_vertices, *start_vertex_labels) =
+            cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices), std::move(*start_vertex_labels));
+        } else {
+          start_vertices =
+            cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices));
+        }
+      }
+
+      //
+      // Need to renumber start_vertices
+      //
+      cugraph::renumber_local_ext_vertices<vertex_t, multi_gpu>(
+        handle_,
+        start_vertices.data(),
+        start_vertices.size(),
+        number_map->data(),
+        graph_view.local_vertex_partition_range_first(),
+        graph_view.local_vertex_partition_range_last(),
+        do_expensive_check_);
+
+      auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
+        cugraph::biased_neighbor_sample(
+          handle_,
+          graph_view,
+          (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+          (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+          (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+          (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
+          raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+          (start_vertex_labels_ != nullptr)
+            ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
+                                                                   start_vertex_labels->size())
+            : std::nullopt,
+          (label_list_ != nullptr)
+            ? std::make_optional(std::make_tuple(
+                raft::device_span<label_t const>{label_list_->as_type<label_t>(),
+                                                 label_list_->size_},
+                raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
+                                                 label_to_comm_rank_->size_}))
+            : std::nullopt,
+          raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
+          rng_state_->rng_state_,
+          options_.return_hops_,
+          options_.with_replacement_,
+          options_.prior_sources_behavior_,
+          options_.dedupe_sources_,
+          do_expensive_check_);
+
+      std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
+
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
+                                                            src.data(),
+                                                            src.size(),
+                                                            number_map->data(),
+                                                            vertex_partition_lasts,
+                                                            do_expensive_check_);
+
+      cugraph::unrenumber_int_vertices<vertex_t, multi_gpu>(handle_,
+                                                            dst.data(),
+                                                            dst.size(),
+                                                            number_map->data(),
+                                                            vertex_partition_lasts,
+                                                            do_expensive_check_);
+
+      std::optional<rmm::device_uvector<vertex_t>> majors{std::nullopt};
+      rmm::device_uvector<vertex_t> minors(0, handle_.get_stream());
+      std::optional<rmm::device_uvector<size_t>> major_offsets{std::nullopt};
+
+      std::optional<rmm::device_uvector<size_t>> label_hop_offsets{std::nullopt};
+
+      std::optional<rmm::device_uvector<vertex_t>> renumber_map{std::nullopt};
+      std::optional<rmm::device_uvector<size_t>> renumber_map_offsets{std::nullopt};
+
+      bool src_is_major = (options_.compression_type_ == cugraph_compression_type_t::CSR) ||
+                          (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+                          (options_.compression_type_ == cugraph_compression_type_t::COO);
+
+      if (options_.renumber_results_) {
+        if (options_.compression_type_ == cugraph_compression_type_t::COO) {
+          // COO
+
+          rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
+          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+          std::tie(output_majors,
+                   minors,
+                   wgt,
+                   edge_id,
+                   edge_type,
+                   label_hop_offsets,
+                   output_renumber_map,
+                   renumber_map_offsets) =
+            cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
+              handle_,
+              std::move(src),
+              std::move(dst),
+              std::move(wgt),
+              std::move(edge_id),
+              std::move(edge_type),
+              std::move(hop),
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
+              offsets ? std::make_optional(
+                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                      : std::nullopt,
+              edge_label ? edge_label->size() : size_t{1},
+              hop ? fan_out_->size_ : size_t{1},
+              src_is_major,
+              do_expensive_check_);
+
+          majors.emplace(std::move(output_majors));
+          renumber_map.emplace(std::move(output_renumber_map));
+        } else {
+          // (D)CSC, (D)CSR
+
+          bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+                                 (options_.compression_type_ == cugraph_compression_type_t::DCSC);
+
+          rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
+          rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
+          std::tie(majors,
+                   output_major_offsets,
+                   minors,
+                   wgt,
+                   edge_id,
+                   edge_type,
+                   label_hop_offsets,
+                   output_renumber_map,
+                   renumber_map_offsets) =
+            cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
+              handle_,
+              std::move(src),
+              std::move(dst),
+              std::move(wgt),
+              std::move(edge_id),
+              std::move(edge_type),
+              std::move(hop),
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<vertex_t const>{
+                    start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
+                : std::nullopt,
+              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                         label_offsets_->as_type<size_t>(), label_offsets_->size_})
+                                     : std::nullopt,
+              offsets ? std::make_optional(
+                          raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                      : std::nullopt,
+              edge_label ? edge_label->size() : size_t{1},
+              hop ? fan_out_->size_ : size_t{1},
+              src_is_major,
+              options_.compress_per_hop_,
+              doubly_compress,
+              do_expensive_check_);
+
+          renumber_map.emplace(std::move(output_renumber_map));
+          major_offsets.emplace(std::move(output_major_offsets));
+        }
+
+        // These are now represented by label_hop_offsets
+        hop.reset();
+        offsets.reset();
+      } else {
+        if (options_.compression_type_ != cugraph_compression_type_t::COO) {
+          CUGRAPH_FAIL("Can only use COO format if not renumbering");
+        }
+
+        std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
+          cugraph::sort_sampled_edgelist(handle_,
+                                         std::move(src),
+                                         std::move(dst),
+                                         std::move(wgt),
+                                         std::move(edge_id),
+                                         std::move(edge_type),
+                                         std::move(hop),
+                                         offsets
+                                           ? std::make_optional(raft::device_span<size_t const>{
+                                               offsets->data(), offsets->size()})
+                                           : std::nullopt,
+                                         edge_label ? edge_label->size() : size_t{1},
+                                         hop ? fan_out_->size_ : size_t{1},
+                                         src_is_major,
+                                         do_expensive_check_);
+
+        majors.emplace(std::move(src));
+        minors = std::move(dst);
+
+        hop.reset();
+        offsets.reset();
+      }
+
+      result_ = new cugraph::c_api::cugraph_sample_result_t{
+        (major_offsets)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*major_offsets, SIZE_T)
+          : nullptr,
+        (majors)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*majors, graph_->vertex_type_)
+          : nullptr,
+        new cugraph::c_api::cugraph_type_erased_device_array_t(minors, graph_->vertex_type_),
+        (edge_id)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*edge_id, graph_->edge_type_)
+          : nullptr,
+        (edge_type) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                        *edge_type, graph_->edge_type_id_type_)
+                    : nullptr,
+        (wgt) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*wgt, graph_->weight_type_)
+              : nullptr,
+        (hop) ? new cugraph::c_api::cugraph_type_erased_device_array_t(*hop, INT32)
+              : nullptr,  // FIXME get rid of this
+        (label_hop_offsets)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(*label_hop_offsets, SIZE_T)
+          : nullptr,
+        (edge_label)
+          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32)
+          : nullptr,
+        (renumber_map) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                           renumber_map.value(), graph_->vertex_type_)
+                       : nullptr,
+        (renumber_map_offsets) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                                   renumber_map_offsets.value(), SIZE_T)
+                               : nullptr};
+    }
+  }
+};
+
 struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   raft::handle_t const& handle_;
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_graph_t* graph_{nullptr};
   cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_offsets_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   int num_edge_types_{};
   cugraph::c_api::cugraph_sampling_options_t options_{};
   bool is_biased_{false};
-  bool is_deprecated_api_{false};
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
@@ -98,16 +788,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_graph_t* graph,
     cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
-    cugraph_type_erased_device_array_view_t const* start_vertex_labels,
     cugraph_type_erased_device_array_view_t const* start_vertex_offsets,
-    cugraph_type_erased_device_array_view_t const* label_list,
     cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
-    cugraph_type_erased_device_array_view_t const* label_offsets,
     cugraph_type_erased_host_array_view_t const* fan_out,
     int num_edge_types,
     cugraph::c_api::cugraph_sampling_options_t options,
     bool is_biased,
-    bool is_deprecated_api,
     bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
@@ -118,26 +804,17 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       start_vertices_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           start_vertices)),
-      start_vertex_labels_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          start_vertex_labels)),
       start_vertex_offsets_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           start_vertex_offsets)),
-      label_list_(reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-        label_list)),
       label_to_comm_rank_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           label_to_comm_rank)),
-      label_offsets_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          label_offsets)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
       num_edge_types_(num_edge_types),
       options_(options),
       is_biased_(is_biased),
-      is_deprecated_api_(is_deprecated_api),
       do_expensive_check_(do_expensive_check)
   {
   }
@@ -196,15 +873,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
       std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
 
-      if (start_vertex_labels_ != nullptr) {
-        start_vertex_labels =
-          rmm::device_uvector<label_t>{start_vertex_labels_->size_, handle_.get_stream()};
-        raft::copy(start_vertex_labels->data(),
-                   start_vertex_labels_->as_type<label_t>(),
-                   start_vertex_labels_->size_,
-                   handle_.get_stream());
-      }
-
       std::optional<rmm::device_uvector<size_t>> start_vertex_offsets{std::nullopt};
 
       if (start_vertex_offsets_ != nullptr) {
@@ -217,15 +885,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       }
 
       if constexpr (multi_gpu) {
-        if (start_vertex_labels) {
-          std::tie(start_vertices, *start_vertex_labels) =
-            cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
-              handle_, std::move(start_vertices), std::move(*start_vertex_labels));
-        } else {
           start_vertices =
             cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
               handle_, std::move(start_vertices));
-        }
       }
 
       //
@@ -249,141 +911,72 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       std::optional<rmm::device_uvector<label_t>> edge_label{std::nullopt};
       std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
 
-      // if 'start_vertex_offsets' is not NULL, leverage the new API
-      if (is_deprecated_api_) {
-        if (is_biased_) {
-          // Call biased neighbor sample
-          auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-            cugraph::biased_neighbor_sample(
-              handle_,
-              graph_view,
-              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-              (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
-              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-              
-              (start_vertex_labels_ != nullptr)
-                ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                      start_vertex_labels->size())
-                : std::nullopt,
-              (label_list_ != nullptr)
-                ? std::make_optional(std::make_tuple(
-                    raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                    label_list_->size_},
-                    raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                    label_to_comm_rank_->size_}))
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              rng_state_->rng_state_,
-              options_.return_hops_,
-              options_.with_replacement_,
+      rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
+
+      cugraph::detail::sequence_fill(handle_.get_stream(),
+                            vertex_type_offsets.begin(),
+                            vertex_type_offsets.size(),
+                            vertex_t{0});
+
+      if (num_edge_types_ > 1) {
+        // call heterogeneous neighbor sample
+        auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+          cugraph::heterogeneous_neighbor_sample(
+            handle_,
+            rng_state_->rng_state_,
+            graph_view,
+            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+            (start_vertex_offsets_ != nullptr)
+              ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                    start_vertex_offsets->size())
+              : std::nullopt,
+            (label_to_comm_rank_ != nullptr)
+              ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                  label_to_comm_rank_->size_})
+              : std::nullopt,
+            raft::host_span<const int>(
+                                      fan_out_->as_type<const int>(), fan_out_->size_),
+            num_edge_types_,
+            cugraph::sampling_flags_t{
               options_.prior_sources_behavior_,
-              options_.dedupe_sources_,
-              do_expensive_check_);
-          
-        } else {
-          // Call uniform neighbor sample
-          auto&& [src, dst, wgt, edge_id, edge_type, hop, edge_label, offsets] =
-            cugraph::uniform_neighbor_sample(
-              handle_,
-              graph_view,
-              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-              
-              (start_vertex_labels_ != nullptr)
-                ? std::make_optional<raft::device_span<label_t const>>(start_vertex_labels->data(),
-                                                                      start_vertex_labels->size())
-                : std::nullopt,
-              (label_list_ != nullptr)
-                ? std::make_optional(std::make_tuple(
-                    raft::device_span<label_t const>{label_list_->as_type<label_t>(),
-                                                    label_list_->size_},
-                    raft::device_span<label_t const>{label_to_comm_rank_->as_type<label_t>(),
-                                                    label_to_comm_rank_->size_}))
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              rng_state_->rng_state_,
               options_.return_hops_,
-              options_.with_replacement_,
+              options_.dedupe_sources_,
+              options_.with_replacement_
+            },
+            do_expensive_check_);
+      } else {
+        // Call homogeneous neighbor sample
+        auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+          cugraph::homogeneous_neighbor_sample(
+            handle_,
+            rng_state_->rng_state_,
+            graph_view,
+            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
+            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+            (start_vertex_offsets_ != nullptr)
+              ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                    start_vertex_offsets->size())
+              : std::nullopt,
+            (label_to_comm_rank_ != nullptr)
+              ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                  label_to_comm_rank_->size_})
+              : std::nullopt,
+            raft::host_span<const int>(
+                                      fan_out_->as_type<const int>(), fan_out_->size_),
+            cugraph::sampling_flags_t{
               options_.prior_sources_behavior_,
+              options_.return_hops_,
               options_.dedupe_sources_,
-              do_expensive_check_);
-        }
-
-      } else { // leverage the new API
-
-        rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
-
-        cugraph::detail::sequence_fill(handle_.get_stream(),
-                              vertex_type_offsets.begin(),
-                              vertex_type_offsets.size(),
-                              vertex_t{0});
-
-        if (true) {
-          // call heterogeneous neighbor sample
-          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
-            cugraph::heterogeneous_neighbor_sample(
-              handle_,
-              rng_state_->rng_state_,
-              graph_view,
-              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-              is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-              (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
-                : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              num_edge_types_,
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
-              do_expensive_check_);
-        } else {
-          // Call homogeneous neighbor sample
-          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
-            cugraph::homogeneous_neighbor_sample(
-              handle_,
-              rng_state_->rng_state_,
-              graph_view,
-              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-              is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-              (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
-                : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
-              do_expensive_check_);
-        }
+              options_.with_replacement_
+            },
+            do_expensive_check_);
       }
 
       std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
@@ -415,6 +1008,14 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                           (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
                           (options_.compression_type_ == cugraph_compression_type_t::COO);
 
+      // Extract the edge_label from the offsets
+      if (offsets) {
+        edge_label = cugraph::c_api::expand_sparse_offsets(
+          raft::device_span<size_t const>{(*offsets).data(), (*offsets).size()},
+          label_t{0},
+          handle_.get_stream());
+      }
+
       if (options_.renumber_results_) {
         if (num_edge_types_ == 1) { // homogeneous renumbering
           if (options_.compression_type_ == cugraph_compression_type_t::COO) {
@@ -442,14 +1043,13 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                options_.retain_seeds_ ? (is_deprecated_api_? std::make_optional(raft::device_span<size_t const>{
-                                          label_offsets_->as_type<size_t>(), label_offsets_->size_}) : std::make_optional(raft::device_span<size_t const>{
-                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_}))
-                                      : std::nullopt,
+                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
+                                          : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
-                edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label ?
+                offsets ? edge_label->size() : size_t{1}, // FIXME: update edge_label ?
                 hop ? fan_out_->size_ : size_t{1},
                 src_is_major,
                 do_expensive_check_);
@@ -487,10 +1087,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                options_.retain_seeds_ ? (is_deprecated_api_? std::make_optional(raft::device_span<size_t const>{
-                                          label_offsets_->as_type<size_t>(), label_offsets_->size_}) : std::make_optional(raft::device_span<size_t const>{
-                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_}))
-                                      : std::nullopt,
+                 options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
+                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
+                                          : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
@@ -1186,22 +1785,17 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
   
   bool is_biased = false;
 
-  neighbor_sampling_functor functor{handle,
-                                    rng_state,
-                                    graph,
-                                    nullptr,
-                                    start_vertices,
-                                    start_vertex_labels, // Deprecated flag
-                                    nullptr, //start_vertex_offsets
-                                    label_list,
-                                    label_to_comm_rank,
-                                    label_offsets,
-                                    fan_out,
-                                    1,
-                                    std::move(options_cpp),
-                                    is_biased,
-                                    true, // is_deprecated_api
-                                    do_expensive_check};
+  uniform_neighbor_sampling_functor functor{handle,
+                                            graph,
+                                            start_vertices,
+                                            start_vertex_labels,
+                                            label_list,
+                                            label_to_comm_rank,
+                                            label_offsets,
+                                            fan_out,
+                                            rng_state,
+                                            std::move(options_cpp),
+                                            do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
@@ -1271,22 +1865,18 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
 
   bool is_biased = true;
 
-  neighbor_sampling_functor functor{handle,
-                                    rng_state,
-                                    graph,
-                                    edge_biases,
-                                    start_vertices,
-                                    start_vertex_labels, // Deprecated flag
-                                    nullptr, //start_vertex_offsets
-                                    label_list,
-                                    label_to_comm_rank,
-                                    label_offsets,
-                                    fan_out,
-                                    1,
-                                    std::move(options_cpp),
-                                    is_biased,
-                                    true, // is_deprecated_api
-                                    do_expensive_check};
+  biased_neighbor_sampling_functor functor{handle,
+                                           graph,
+                                           edge_biases,
+                                           start_vertices,
+                                           start_vertex_labels,
+                                           label_list,
+                                           label_to_comm_rank,
+                                           label_offsets,
+                                           fan_out,
+                                           rng_state,
+                                           std::move(options_cpp),
+                                          do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
@@ -1356,16 +1946,12 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    nullptr, // start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
-                                    nullptr, // label_list /*Deprecated flag*/
                                     label_to_comm_rank,
-                                    nullptr, // label_offsets /*Deprecated flag*/
                                     fan_out,
                                     num_edge_types,
                                     std::move(options_cpp),
                                     is_biased,
-                                    false, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
@@ -1436,16 +2022,12 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                                     graph,
                                     edge_biases,
                                     start_vertices,
-                                    nullptr, //start_vertex_labels /*Deprecated flag*/
                                     start_vertex_offsets,
-                                    nullptr,//label_list /*Deprecated flag*/
                                     label_to_comm_rank,
-                                    nullptr, //label_offsets /*Deprecated flag*/
                                     fan_out,
-                                    1,
+                                    1, // num_edge_types
                                     std::move(options_cpp),
                                     is_biased,
-                                    false, // is_deprecated_api
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index cd62bbd06c5..5bc8db687a7 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -75,6 +75,28 @@ neighbor_sample_impl(raft::handle_t const& handle,
   std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
   graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
 
+  if constexpr (!multi_gpu) {
+    CUGRAPH_EXPECTS(!label_to_output_comm_rank,
+                    "cannot specify output GPU mapping in SG implementation");
+  }
+
+  CUGRAPH_EXPECTS(
+    !label_to_output_comm_rank || starting_vertex_labels,
+    "cannot specify output GPU mapping without also specifying starting_vertex_labels");
+
+  if (do_expensive_check) {
+    if (edge_bias_view) {
+      auto [num_negative_edge_weights, num_overflows] =
+        check_edge_bias_values(handle, graph_view, *edge_bias_view);
+
+      CUGRAPH_EXPECTS(
+        num_negative_edge_weights == 0,
+        "Invalid input argument: input edge bias values should have non-negative values.");
+      CUGRAPH_EXPECTS(num_overflows == 0,
+                      "Invalid input argument: sum of neighboring edge bias values should not "
+                      "exceed std::numeric_limits<bias_t>::max() for any vertex.");
+    }
+  }
 
   CUGRAPH_EXPECTS(fan_out.size() > 0, "Invalid input argument: number of levels must be non-zero.");
   CUGRAPH_EXPECTS(
@@ -89,8 +111,6 @@ neighbor_sample_impl(raft::handle_t const& handle,
     
     cugraph::fill_edge_property(handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
 
-    modified_graph_view.attach_edge_mask(edge_mask.view());
-
     cugraph::transform_e(
       handle,
       modified_graph_view,
@@ -106,29 +126,6 @@ neighbor_sample_impl(raft::handle_t const& handle,
     edge_masks_vector.push_back(std::move(edge_mask));
   }
 
-  if constexpr (!multi_gpu) {
-    CUGRAPH_EXPECTS(!label_to_output_comm_rank,
-                    "cannot specify output GPU mapping in SG implementation");
-  }
-
-  CUGRAPH_EXPECTS(
-    !label_to_output_comm_rank || starting_vertex_labels,
-    "cannot specify output GPU mapping without also specifying starting_vertex_labels");
-
-  if (do_expensive_check) {
-    if (edge_bias_view) {
-      auto [num_negative_edge_weights, num_overflows] =
-        check_edge_bias_values(handle, graph_view, *edge_bias_view);
-
-      CUGRAPH_EXPECTS(
-        num_negative_edge_weights == 0,
-        "Invalid input argument: input edge bias values should have non-negative values.");
-      CUGRAPH_EXPECTS(num_overflows == 0,
-                      "Invalid input argument: sum of neighboring edge bias values should not "
-                      "exceed std::numeric_limits<bias_t>::max() for any vertex.");
-    }
-  }
-
   std::vector<rmm::device_uvector<vertex_t>> level_result_src_vectors{};
   std::vector<rmm::device_uvector<vertex_t>> level_result_dst_vectors{};
   auto level_result_weight_vectors =
@@ -171,9 +168,6 @@ neighbor_sample_impl(raft::handle_t const& handle,
   }
 
   std::vector<size_t> level_sizes{};
-  
-  auto cur_graph_view = modified_graph_view;
-
 
   // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1
   auto num_hops = (fan_out.size() % num_edge_types) ? (fan_out.size() / num_edge_types) : (fan_out.size() / num_edge_types + 1);
@@ -188,10 +182,12 @@ neighbor_sample_impl(raft::handle_t const& handle,
       std::optional<rmm::device_uvector<edge_type_t>> edge_types{std::nullopt};
       std::optional<rmm::device_uvector<int32_t>> labels{std::nullopt};
 
+      modified_graph_view.attach_edge_mask(edge_masks_vector[edge_type_id].view());
+
       if (k_level > 0) {
         std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
           sample_edges(handle,
-                      cur_graph_view,
+                      modified_graph_view,
                       edge_weight_view,
                       edge_id_view,
                       edge_type_view,
@@ -204,7 +200,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
       } else {
         std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
           gather_one_hop_edgelist(handle,
-                                  cur_graph_view,
+                                  modified_graph_view,
                                   edge_weight_view,
                                   edge_id_view,
                                   edge_type_view,
@@ -224,7 +220,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
 
     // FIXME:  We should modify vertex_partition_range_lasts to return a raft::host_span
     //  rather than making a copy.
-    auto vertex_partition_range_lasts = cur_graph_view.vertex_partition_range_lasts();
+    auto vertex_partition_range_lasts = modified_graph_view.vertex_partition_range_lasts();
     std::tie(frontier_vertices, frontier_vertex_labels, vertex_used_as_source) =
       prepare_next_frontier(
         handle,
@@ -237,7 +233,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
                                     level_result_label_vectors->back().size()))
                                 : std::nullopt,
         std::move(vertex_used_as_source),
-        cur_graph_view.local_vertex_partition_view(),
+        modified_graph_view.local_vertex_partition_view(),
         vertex_partition_range_lasts,
         prior_sources_behavior,
         dedupe_sources,
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 8addad740e9..9366a31605f 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -75,7 +75,7 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_type_erased_device_array_view_t* label_list,
         const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
         const cugraph_type_erased_device_array_view_t* label_offsets,
-        const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout,
+        const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out,
         const cugraph_sampling_options_t* options,
         bool_t is_biased,
         bool_t do_expensive_check,
@@ -151,13 +151,13 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
             cugraph_graph_t* graph,
             const cugraph_type_erased_host_array_view_t* edge_type_offsets,
             const cugraph_type_erased_host_array_view_t* fanout,
-            cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fanout,
+            cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
             cugraph_error_t** error
         )
 
     cdef void \
         cugraph_heterogeneous_fan_out_free(
-            cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fanout);
+            cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out);
     # negative sampling
     cdef cugraph_error_code_t \
         cugraph_negative_sampling(

From 4857b36635cfcb9fce9465c0cd301c1a9637e00f Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 23 Sep 2024 11:48:41 -0700
Subject: [PATCH 046/117] call scatter instead of gather and fix type bug

---
 cpp/src/sampling/detail/conversion_utilities.cu | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index 5e6f4b00fd2..c6e396b7024 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -35,7 +35,7 @@ rmm::device_uvector<int32_t> flatten_label_map(
 {
   rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
 
-  label_t max_label = thrust::reduce(handle.get_thrust_policy(),
+  label_t max_label = thrust::scatter(handle.get_thrust_policy(),
                                      std::get<0>(label_to_output_comm_rank).begin(),
                                      std::get<0>(label_to_output_comm_rank).end(),
                                      label_t{0},
@@ -43,7 +43,7 @@ rmm::device_uvector<int32_t> flatten_label_map(
 
   label_map.resize(max_label, handle.get_stream());
 
-  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), label_t{0});
+  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t);
   thrust::gather(handle.get_thrust_policy(),
                  std::get<0>(label_to_output_comm_rank).begin(),
                  std::get<0>(label_to_output_comm_rank).end(),

From 263b6ac026bee400b32149b0d1a4c5af575412cc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 23 Sep 2024 11:55:32 -0700
Subject: [PATCH 047/117] fix typo

---
 cpp/src/sampling/detail/conversion_utilities.cu | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index c6e396b7024..bc6cb128561 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -35,7 +35,7 @@ rmm::device_uvector<int32_t> flatten_label_map(
 {
   rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
 
-  label_t max_label = thrust::scatter(handle.get_thrust_policy(),
+  label_t max_label = thrust::reduce(handle.get_thrust_policy(),
                                      std::get<0>(label_to_output_comm_rank).begin(),
                                      std::get<0>(label_to_output_comm_rank).end(),
                                      label_t{0},
@@ -43,12 +43,12 @@ rmm::device_uvector<int32_t> flatten_label_map(
 
   label_map.resize(max_label, handle.get_stream());
 
-  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t);
-  thrust::gather(handle.get_thrust_policy(),
-                 std::get<0>(label_to_output_comm_rank).begin(),
-                 std::get<0>(label_to_output_comm_rank).end(),
-                 std::get<1>(label_to_output_comm_rank).begin(),
-                 label_map.begin());
+  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
+  thrust::scatter(handle.get_thrust_policy(),
+                  std::get<0>(label_to_output_comm_rank).begin(),
+                  std::get<0>(label_to_output_comm_rank).end(),
+                  std::get<1>(label_to_output_comm_rank).begin(),
+                  label_map.begin());
 
   return label_map;
 }

From 9dff3ab16a557f370d7a54deae5d21fc317cbe7a Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 24 Sep 2024 10:25:26 -0700
Subject: [PATCH 048/117] update neighbor sample API

---
 cpp/include/cugraph/sampling_functions.hpp    | 194 +++++++++-
 cpp/include/cugraph_c/sampling_algorithms.h   | 127 ++++++-
 cpp/src/c_api/neighbor_sampling.cpp           | 348 ++++++++++++++----
 cpp/src/sampling/neighbor_sampling_impl.hpp   | 145 +++++++-
 .../sampling/neighbor_sampling_mg_v32_e32.cu  |  14 +-
 .../sampling/neighbor_sampling_mg_v32_e64.cu  |  14 +-
 .../sampling/neighbor_sampling_mg_v64_e64.cu  |  14 +-
 .../sampling/neighbor_sampling_sg_v32_e32.cu  |  14 +-
 .../sampling/neighbor_sampling_sg_v32_e64.cu  |  14 +-
 .../sampling/neighbor_sampling_sg_v64_e64.cu  |  14 +-
 10 files changed, 741 insertions(+), 157 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 1f82ef67127..822a6a2d991 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -271,11 +271,12 @@ struct sampling_flags_t {
 };
 
 /**
- * @brief Heterogeneous Neighborhood Sampling.
+ * @brief Heterogeneous Uniform Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
- * The branching out to select outgoing neighbors is performed with homogeneous fanouts.
+ * randomly selects (uniformly) from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with heterogeneous fanouts
+ * where the number of edge types is bigger than 1.
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
  * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
@@ -309,11 +310,94 @@ struct sampling_flags_t {
  * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
  * @param edge_id_view Optional view object holding edge ids for @p graph_view.
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+ * @param starting_vertex_offsets Optional device span of offsets identifying the range of
+ * starting vertex values for this label.
+ * @param label_to_output_comm_rank Optional device span identifying which rank should get each
+ * vertex label.  This should be the same on each rank.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types + edge_type_id]'
+ * @param num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
+ * sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
+ * @param flags A set of flags indicating which sampling features should be used.
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  edge_type_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check = false);
+
+/**
+ * @brief Heterogeneous Biased Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with heterogeneous fanouts
+ * where the number of edge types is bigger than 1.
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The offsets array (optional) identifies the offset for each label.
+ *
+ * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * in the return will be a CSR-style offsets array to identify the beginning of each label range in
+ * the output vectors.
+ *
+ * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
+ * for a particular label are returned on the specified rank.  This will result in the offsets array
+ * on other GPUs indicating that there are no entries for that label (`offsets[i] == offsets[i+1]`).
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam bias_t Type of bias. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
  * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
  * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
- * uniform sampling.
+ * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
@@ -344,14 +428,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -361,10 +445,10 @@ heterogeneous_neighbor_sample(
   bool do_expensive_check = false);
 
 /**
- * @brief Homogeneous Neighborhood Sampling.
+ * @brief Homogeneous Uniform Neighborhood Sampling.
  *
  * This function traverses from a set of starting vertices, traversing outgoing edges and
- * randomly selects (with edge biases or not) from these outgoing neighbors to extract a subgraph.
+ * randomly selects (uniformly) from these outgoing neighbors to extract a subgraph.
  * The branching out to select outgoing neighbors is performed with homogeneous fanouts
  *
  * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
@@ -399,11 +483,91 @@ heterogeneous_neighbor_sample(
  * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
  * @param edge_id_view Optional view object holding edge ids for @p graph_view.
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
- * @param edge_bias_view Optional view object holding edge biases (to be used in biased sampling) for @p
+ * @param starting_vertices Device span of starting vertex IDs for the sampling.
+ * In a multi-gpu context the starting vertices should be local to this GPU.
+  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
+ * starting vertex values for this label.
+ * @param label_to_output_comm_rank Optional device span identifying which rank should get each
+ * vertex label.  This should be the same on each rank.
+ * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
+ * level.
+ * @param flags A set of flags indicating which sampling features should be used.
+ * @param do_expensive_check A flag to run expensive checks for input arguments (if set to `true`).
+ * @return tuple device vectors (vertex_t source_vertex, vertex_t destination_vertex,
+ * optional weight_t weight, optional edge_t edge id, optional edge_type_t edge type,
+ * optional int32_t hop, optional label_t label, optional size_t offsets)
+ */
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check = false);
+
+/**
+ * @brief Homogeneous Biased Neighborhood Sampling.
+ *
+ * This function traverses from a set of starting vertices, traversing outgoing edges and
+ * randomly selects (with edge biases) from these outgoing neighbors to extract a subgraph.
+ * The branching out to select outgoing neighbors is performed with homogeneous fanouts
+ *
+ * Output from this function is a tuple of vectors (src, dst, weight, edge_id, edge_type, hop,
+ * offsets), identifying the randomly selected edges.  src is the source vertex, dst is the
+ * destination vertex, weight (optional) is the edge weight, edge_id (optional) identifies the edge
+ * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
+ * encountered in.  The offsets array (optional) identifies the offset for each label.
+ *
+ * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * offsets values in the return set will be std::nullopt.
+ *
+ * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * in the return will be a CSR-style offsets array to identify the beginning of each label range in
+ * the output vectors.
+ *
+ * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
+ * for a particular label are returned on the specified rank.  This will result in the offsets array
+ * on other GPUs indicating that there are no entries for that label (`offsets[i] == offsets[i+1]`).
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam edge_t Type of edge identifiers. Needs to be an integral type.
+ * @tparam weight_t Type of edge weights. Needs to be a floating point type.
+ * @tparam edge_type_t Type of edge type. Needs to be an integral type.
+ * @tparam bias_t Type of bias. Needs to be an integral type.
+ * @tparam store_transposed Flag indicating whether sources (if false) or destinations (if
+ * true) are major indices
+ * @tparam multi_gpu Flag indicating whether template instantiation should target single-GPU (false)
+ * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
+ * handles to various CUDA libraries) to run graph algorithms.
+ * @param rng_state A pre-initialized raft::RngState object for generating random numbers
+ * @param graph_view Graph View object to generate NBR Sampling on.
+ * @param edge_weight_view Optional view object holding edge weights for @p graph_view.
+ * @param edge_id_view Optional view object holding edge ids for @p graph_view.
+ * @param edge_type_view Optional view object holding edge types for @p graph_view.
+ * @param edge_bias_view View object holding edge biases (to be used in biased sampling) for @p
  * graph_view. Bias values should be non-negative and the sum of edge bias values from any vertex
  * should not exceed std::numeric_limits<bias_t>::max(). 0 bias value indicates that the
- * corresponding edge can never be selected. passing std::nullopt as the edge biases will result in
- * uniform sampling.
+ * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
   * @param starting_vertex_offsets Optional device span of offsets identifying the range of
@@ -433,14 +597,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 1b2d15dea12..6827f23e986 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -491,10 +491,10 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_error_t** error);
 
 /**
- * @brief     Heterogeneous Neighborhood Sampling
+ * @brief     Heterogeneous Uniform Neighborhood Sampling
  *
- * Returns a sample of the neighborhood around specified start vertices with edge biases or not
- * and homogeneous fan_out types.
+ * Returns a sample of the neighborhood around specified start vertices and heterogeneous
+ * fan_out types. The neighborhood is sampled uniformly.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -506,8 +506,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  *  * @param [in,out] rng_state State of the random number generator, updated with each call
  * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
  *                           needs to be transposed
- * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
- * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
@@ -534,7 +532,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -550,11 +548,12 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
+
 /**
- * @brief     Homogeneous Neighborhood Sampling
+ * @brief     Heterogeneous Biased Neighborhood Sampling
  *
- * Returns a sample of the neighborhood around specified start vertices with edge biases or not
- * and homogeneous fan_out types.
+ * Returns a sample of the neighborhood around specified start vertices and heterogeneous
+ * fan_out types. The neighborhood is sampled with biases.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -579,6 +578,9 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
+ *                            sample whereas a value greater than 1 translates to heterogeneous neighbor
+ *                            sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  is_biased
@@ -591,7 +593,7 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -599,13 +601,116 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sample_heterogeneous_fan_out_t* fan_out,
+  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
+/**
+ * @brief     Homogeneous Uniform Neighborhood Sampling
+ *
+ * Returns a sample of the neighborhood around specified start vertices and heterogeneous
+ * fan_out types. The neighborhood is sampled uniformly.
+ * Optionally, each start vertex can be associated with a label, allowing the caller to specify
+ * multiple batches of sampling requests in the same function call - which should improve GPU
+ * utilization.
+ *
+ * If label is NULL then all start vertices will be considered part of the same batch and the
+ * return value will not have a label column.
+ *
+ * @param [in]  handle       Handle for accessing resources
+ *  * @param [in,out] rng_state State of the random number generator, updated with each call
+ * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
+ *                           needs to be transposed
+ * @param [in]  start_vertices Device array of start vertices for the sampling
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
+ * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
+ * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
+ * If not specifed the data is not organized in output.  If specified then the all data in the range
+ * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_offsets is also specified.
+ * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                            We only support fan_out values of type INT32
+ * @param [in]  sampling_options
+ *                           Opaque pointer defining the sampling options.
+ * @param [in]  do_expensive_check
+ *                           A flag to run expensive checks for input arguments (if set to true)
+ * @param [out]  result      Output from the uniform_neighbor_sample call
+ * @param [out] error        Pointer to an error object storing details of any error.  Will
+ *                           be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error);
+
+/**
+ * @brief     Homogeneous Biased Neighborhood Sampling
+ *
+ * Returns a sample of the neighborhood around specified start vertices and heterogeneous
+ * fan_out types. The neighborhood is sampled uniformly.
+ * Optionally, each start vertex can be associated with a label, allowing the caller to specify
+ * multiple batches of sampling requests in the same function call - which should improve GPU
+ * utilization.
+ *
+ * If label is NULL then all start vertices will be considered part of the same batch and the
+ * return value will not have a label column.
+ *
+ * @param [in]  handle       Handle for accessing resources
+ *  * @param [in,out] rng_state State of the random number generator, updated with each call
+ * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
+ *                           needs to be transposed
+ * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
+ * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
+ * @param [in]  start_vertices Device array of start vertices for the sampling
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
+ *                            parameter is only used with the retain_seeds option.
+ * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
+ * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
+ * If not specifed the data is not organized in output.  If specified then the all data in the range
+ * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
+ * cannot be specified unless @p start_vertex_offsets is also specified.
+ * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
+ *                            We only support fan_out values of type INT32
+ * @param [in]  sampling_options
+ *                           Opaque pointer defining the sampling options.
+ * @param [in]  do_expensive_check
+ *                           A flag to run expensive checks for input arguments (if set to true)
+ * @param [out]  result      Output from the uniform_neighbor_sample call
+ * @param [out] error        Pointer to an error object storing details of any error.  Will
+ *                           be populated if error code is not CUGRAPH_SUCCESS
+ * @return error code
+ */
+cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_edge_property_view_t* edge_biases,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error);
+
+
 /**
  * @deprecated This call should be replaced with cugraph_sample_result_get_majors
  * @brief     Get the source vertices from the sampling algorithm result
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index a084e0fb5d7..04090b22a1d 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -917,66 +917,131 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                             vertex_type_offsets.begin(),
                             vertex_type_offsets.size(),
                             vertex_t{0});
+      
+      // FIXME: For biased sampling, the user should pass either biases or edge weights,
+      // otherwised throw an error and suggest them to use uniform neighbor sample instead
 
       if (num_edge_types_ > 1) {
         // call heterogeneous neighbor sample
-        auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
-          cugraph::heterogeneous_neighbor_sample(
-            handle_,
-            rng_state_->rng_state_,
-            graph_view,
-            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-            (start_vertex_offsets_ != nullptr)
-              ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                    start_vertex_offsets->size())
-              : std::nullopt,
-            (label_to_comm_rank_ != nullptr)
-              ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                  label_to_comm_rank_->size_})
-              : std::nullopt,
-            raft::host_span<const int>(
-                                      fan_out_->as_type<const int>(), fan_out_->size_),
-            num_edge_types_,
-            cugraph::sampling_flags_t{
-              options_.prior_sources_behavior_,
-              options_.return_hops_,
-              options_.dedupe_sources_,
-              options_.with_replacement_
-            },
-            do_expensive_check_);
+        if (is_biased_) {
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::heterogeneous_biased_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              num_edge_types_,
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+        } else {
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::heterogeneous_uniform_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              num_edge_types_,
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+
+        }
       } else {
         // Call homogeneous neighbor sample
-        auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
-          cugraph::homogeneous_neighbor_sample(
-            handle_,
-            rng_state_->rng_state_,
-            graph_view,
-            (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-            (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-            (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-            is_biased_ ? ((edge_biases != nullptr) ? std::make_optional(*edge_biases) : std::make_optional(edge_weights->view())) : std::nullopt,
-            raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-            (start_vertex_offsets_ != nullptr)
-              ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                    start_vertex_offsets->size())
-              : std::nullopt,
-            (label_to_comm_rank_ != nullptr)
-              ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                  label_to_comm_rank_->size_})
-              : std::nullopt,
-            raft::host_span<const int>(
-                                      fan_out_->as_type<const int>(), fan_out_->size_),
-            cugraph::sampling_flags_t{
-              options_.prior_sources_behavior_,
-              options_.return_hops_,
-              options_.dedupe_sources_,
-              options_.with_replacement_
-            },
-            do_expensive_check_);
+        if (is_biased_) {
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::homogeneous_biased_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+        } else {
+          if (is_biased_) {
+          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+            cugraph::homogeneous_uniform_neighbor_sample(
+              handle_,
+              rng_state_->rng_state_,
+              graph_view,
+              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+              (start_vertex_offsets_ != nullptr)
+                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
+                                                                      start_vertex_offsets->size())
+                : std::nullopt,
+              (label_to_comm_rank_ != nullptr)
+                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
+                                                    label_to_comm_rank_->size_})
+                : std::nullopt,
+              raft::host_span<const int>(
+                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              cugraph::sampling_flags_t{
+                options_.prior_sources_behavior_,
+                options_.return_hops_,
+                options_.dedupe_sources_,
+                options_.with_replacement_
+              },
+              do_expensive_check_);
+          }
+        }
       }
 
       std::vector<vertex_t> vertex_partition_lasts = graph_view.vertex_partition_range_lasts();
@@ -1880,7 +1945,72 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  int num_edge_types,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error)
+{
+  auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
+  
+  // FIXME: Should we maintain this contition?
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "must specify start_vertex_offsets if retain_seeds is true",
+               *error);
+  
+  CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
+                 (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                    start_vertex_offsets)
+                    ->type_ == SIZE_T),
+               CUGRAPH_INVALID_INPUT,
+               "start_vertex_offsets should be of type size_t",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
+               *error);
+    
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out should be of type int",
+    *error);
+
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
+                 reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                   start_vertices)
+                   ->type_,
+               CUGRAPH_INVALID_INPUT,
+               "vertex type of graph and start_vertices must match",
+               *error);
+
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    nullptr,
+                                    start_vertices,
+                                    start_vertex_offsets,
+                                    label_to_comm_rank,
+                                    fan_out,
+                                    num_edge_types,
+                                    std::move(options_cpp),
+                                    FALSE,
+                                    do_expensive_check};
+  return cugraph::c_api::run_algorithm(graph, functor, result, error);
+}
+
+cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -1891,21 +2021,18 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
-  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
-  if (is_biased) {
-    CAPI_EXPECTS(
-      (edge_biases != nullptr) ||
-        (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
-      CUGRAPH_INVALID_INPUT,
-      "edge_biases is required if the graph is not weighted",
-      *error);
-  }
+  CAPI_EXPECTS(
+    (edge_biases != nullptr) ||
+      (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
+    CUGRAPH_INVALID_INPUT,
+    "edge_biases is required if the graph is not weighted",
+    *error);
   
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
@@ -1951,12 +2078,77 @@ cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
                                     fan_out,
                                     num_edge_types,
                                     std::move(options_cpp),
-                                    is_biased,
+                                    TRUE,
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
-cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
+cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
+  const cugraph_resource_handle_t* handle,
+  cugraph_rng_state_t* rng_state,
+  cugraph_graph_t* graph,
+  const cugraph_type_erased_device_array_view_t* start_vertices,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
+  const cugraph_type_erased_host_array_view_t* fan_out,
+  const cugraph_sampling_options_t* options,
+  bool_t do_expensive_check,
+  cugraph_sample_result_t** result,
+  cugraph_error_t** error)
+{
+  auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
+  
+  // FIXME: Should we maintain this contition?
+  CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "must specify start_vertex_offsets if retain_seeds is true",
+               *error);
+  
+  CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
+                 (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                    start_vertex_offsets)
+                    ->type_ == SIZE_T),
+               CUGRAPH_INVALID_INPUT,
+               "start_vertex_offsets should be of type size_t",
+               *error);
+
+  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
+               CUGRAPH_INVALID_INPUT,
+               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
+               *error);
+  
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
+                fan_out)
+                ->type_ == INT32,
+            CUGRAPH_INVALID_INPUT,
+            "fan_out type must be INT32",
+            *error);
+    
+  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
+                 reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
+                   start_vertices)
+                   ->type_,
+               CUGRAPH_INVALID_INPUT,
+               "vertex type of graph and start_vertices must match",
+               *error);
+
+
+  neighbor_sampling_functor functor{handle,
+                                    rng_state,
+                                    graph,
+                                    nullptr,
+                                    start_vertices,
+                                    start_vertex_offsets,
+                                    label_to_comm_rank,
+                                    fan_out,
+                                    1, // num_edge_types
+                                    std::move(options_cpp),
+                                    FALSE,
+                                    do_expensive_check};
+  return cugraph::c_api::run_algorithm(graph, functor, result, error);
+}
+
+cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -1966,21 +2158,18 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
-  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
-  if (is_biased) {
-    CAPI_EXPECTS(
-      (edge_biases != nullptr) ||
-        (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
-      CUGRAPH_INVALID_INPUT,
-      "edge_biases is required if the graph is not weighted",
-      *error);
-  }
+  CAPI_EXPECTS(
+    (edge_biases != nullptr) ||
+      (reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->edge_weights_ != nullptr),
+    CUGRAPH_INVALID_INPUT,
+    "edge_biases is required if the graph is not weighted",
+    *error);
   
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
@@ -2027,11 +2216,16 @@ cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
                                     fan_out,
                                     1, // num_edge_types
                                     std::move(options_cpp),
-                                    is_biased,
+                                    TRUE,
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
+
+
+
+
+
 extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
   const cugraph_resource_handle_t* handle,
   cugraph_graph_t* graph,
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 5bc8db687a7..899ff644179 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -490,6 +490,73 @@ biased_neighbor_sample(
     do_expensive_check);
 }
 
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  edge_type_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check)
+{ 
+  using bias_t = weight_t; // dummy
+  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
+
+  if (starting_vertex_offsets)
+    starting_vertex_labels =
+      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
+  
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
+    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        edge_id_view,
+        edge_type_view,
+        std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
+        starting_vertices,
+        starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
+                                    starting_vertex_labels.data(), starting_vertex_labels.size()})
+                                : std::nullopt,
+        label_to_output_comm_rank,
+        fan_out,
+        num_edge_types,
+        sampling_flags.return_hops,
+        sampling_flags.with_replacement,
+        sampling_flags.prior_sources_behavior,
+        sampling_flags.dedupe_sources,
+        do_expensive_check);
+  
+  return std::make_tuple(std::move(majors),
+                         std::move(minors),
+                         std::move(weights),
+                         std::move(edge_ids),
+                         std::move(edge_types),
+                         std::move(hops),
+                         std::move(offsets));
+  
+}
+
 template <typename vertex_t,
           typename edge_t,
           typename weight_t,
@@ -504,14 +571,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -534,7 +601,7 @@ heterogeneous_neighbor_sample(
         edge_weight_view,
         edge_id_view,
         edge_type_view,
-        edge_bias_view,
+        std::make_optional(edge_bias_view),
         starting_vertices,
         starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
                                     starting_vertex_labels.data(), starting_vertex_labels.size()})
@@ -555,7 +622,73 @@ heterogeneous_neighbor_sample(
                          std::move(edge_types),
                          std::move(hops),
                          std::move(offsets));
+}
+
+template <typename vertex_t,
+          typename edge_t,
+          typename weight_t,
+          typename edge_type_t,
+          bool store_transposed,
+          bool multi_gpu>
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<vertex_t>,
+           std::optional<rmm::device_uvector<weight_t>>,
+           std::optional<rmm::device_uvector<edge_t>>,
+           std::optional<rmm::device_uvector<edge_type_t>>,
+           std::optional<rmm::device_uvector<int32_t>>,
+           std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
+  std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
+  std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
+  raft::device_span<vertex_t const> starting_vertices,
+  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  // edge_type_t num_edge_types, /*argument not needed for homogeneous neighbor sample*/
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check)
+{
+  using bias_t = weight_t; // dummy
+  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
+
+  if (starting_vertex_offsets)
+    starting_vertex_labels =
+      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
+
+  
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
+    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+    handle,
+    rng_state,
+    graph_view,
+    edge_weight_view,
+    edge_id_view,
+    edge_type_view,
+    std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
+    starting_vertices,
+    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
+                                starting_vertex_labels.data(), starting_vertex_labels.size()})
+                            : std::nullopt,
+    label_to_output_comm_rank,
+    fan_out,
+    edge_type_t{1},
+    sampling_flags.return_hops,
+    sampling_flags.with_replacement,
+    sampling_flags.prior_sources_behavior,
+    sampling_flags.dedupe_sources,
+    do_expensive_check);
   
+  return std::make_tuple(std::move(majors),
+                         std::move(minors),
+                         std::move(weights),
+                         std::move(edge_ids),
+                         std::move(edge_types),
+                         std::move(hops),
+                         std::move(offsets));
 }
 
 template <typename vertex_t,
@@ -572,14 +705,14 @@ std::tuple<rmm::device_uvector<vertex_t>,
            std::optional<rmm::device_uvector<edge_type_t>>,
            std::optional<rmm::device_uvector<int32_t>>,
            std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu> const& graph_view,
   std::optional<edge_property_view_t<edge_t, weight_t const*>> edge_weight_view,
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<edge_t, bias_t const*>> edge_bias_view,
+  edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -604,7 +737,7 @@ homogeneous_neighbor_sample(
     edge_weight_view,
     edge_id_view,
     edge_type_view,
-    edge_bias_view,
+    std::make_optional(edge_bias_view),
     starting_vertices,
     starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
                                 starting_vertex_labels.data(), starting_vertex_labels.size()})
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
index 99b0ea82ea5..6384f80f8e8 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
@@ -134,14 +134,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -157,14 +156,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -180,14 +179,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -202,14 +200,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
index 159c2b3c372..0b153b976c4 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
@@ -135,14 +135,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -158,14 +157,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -181,14 +180,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -203,14 +201,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
index 0d69815778f..da2c89dbeb7 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
@@ -134,14 +134,13 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -157,14 +156,14 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -180,14 +179,13 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -202,14 +200,14 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, true> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
index 25a597b53fd..d3f7cf3ce74 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
@@ -134,14 +134,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -157,14 +156,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -180,14 +179,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<int32_t>>,
                   std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -202,14 +200,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int32_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int32_t, double const*>> edge_bias_view,
+  edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
index feb1c3cf052..205749bb9fc 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
@@ -134,14 +134,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -157,14 +156,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -180,14 +179,13 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -202,14 +200,14 @@ template std::tuple<rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int32_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
index 334fef6ff8f..97cccb9d307 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
@@ -134,14 +134,13 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -157,14 +156,14 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_neighbor_sample(
+heterogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -180,14 +179,13 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
@@ -202,14 +200,14 @@ template std::tuple<rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<int32_t>>,
                     std::optional<rmm::device_uvector<size_t>>>
-homogeneous_neighbor_sample(
+homogeneous_biased_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
   graph_view_t<int64_t, int64_t, false, false> const& graph_view,
   std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_bias_view,
+  edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
   std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,

From 33c8b3d0906bd74d09eea50c94afd5798fb57106 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 25 Sep 2024 14:08:22 -0700
Subject: [PATCH 049/117] update CAPI

---
 .../cugraph/detail/shuffle_wrappers.hpp       |  23 ++++
 .../cugraph/detail/utility_wrappers.hpp       |  16 +++
 cpp/include/cugraph/sampling_functions.hpp    |  28 ++--
 cpp/include/cugraph_c/sampling_algorithms.h   |  28 ----
 cpp/src/c_api/neighbor_sampling.cpp           | 121 +++++++++---------
 cpp/src/detail/utility_wrappers_32.cu         |  10 ++
 cpp/src/detail/utility_wrappers_64.cu         |  10 ++
 cpp/src/detail/utility_wrappers_impl.cuh      |  16 +++
 cpp/src/sampling/neighbor_sampling_impl.hpp   |  49 ++-----
 .../sampling/neighbor_sampling_mg_v32_e32.cu  |   8 +-
 .../sampling/neighbor_sampling_mg_v32_e64.cu  |   8 +-
 .../sampling/neighbor_sampling_mg_v64_e64.cu  |   8 +-
 .../sampling/neighbor_sampling_sg_v32_e32.cu  |   8 +-
 .../sampling/neighbor_sampling_sg_v32_e64.cu  |   8 +-
 .../sampling/neighbor_sampling_sg_v64_e64.cu  |   8 +-
 cpp/src/utilities/shuffle_vertices.cuh        |  76 +++++++++--
 .../utilities/shuffle_vertices_mg_v32_fp.cu   |  14 ++
 .../shuffle_vertices_mg_v32_integral.cu       |  15 +++
 .../utilities/shuffle_vertices_mg_v64_fp.cu   |  14 ++
 .../shuffle_vertices_mg_v64_integral.cu       |  14 ++
 .../pylibcugraph/pylibcugraph/CMakeLists.txt  |   2 -
 21 files changed, 306 insertions(+), 178 deletions(-)

diff --git a/cpp/include/cugraph/detail/shuffle_wrappers.hpp b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
index 7dffcce298a..68ebdf01296 100644
--- a/cpp/include/cugraph/detail/shuffle_wrappers.hpp
+++ b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
@@ -142,6 +142,29 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<vertex_t>&& vertices,
   rmm::device_uvector<value_t>&& values);
 
+/**
+ * @brief Shuffle external (i.e. before renumbering) vertex & values pairs to their local GPU based
+ * on vertex partitioning.
+ *
+ * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
+ * @tparam value_t Type of values.
+ *
+ * @param[in] handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
+ * and handles to various CUDA libraries) to run graph algorithms.
+ * @param[in] vertices Vertices to shuffle.
+ * @param[in] values_0 First values to shuffle.
+ * @param[in] values_1 Second values to shuffle.
+ *
+ * @return Tuple of vectors storing shuffled vertex & value pairs.
+ */
+template <typename vertex_t, typename value0_t, typename value1_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<vertex_t>&& vertices,
+  rmm::device_uvector<value0_t>&& values_0,
+  rmm::device_uvector<value1_t>&& values_1);
+
 /**
  * @brief Permute a range.
  *
diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index 3d99b85556b..fa61c91a4ab 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -65,6 +65,22 @@ void uniform_random_fill(rmm::cuda_stream_view const& stream_view,
 template <typename value_t>
 void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, value_t value);
 
+
+/**
+ * @brief    Increment the values of a buffer by a constant value
+ *
+ * @tparam      value_t      type of the value to operate on
+ *
+ * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
+ * and handles to various CUDA libraries) to run graph algorithms.
+ * @param[out]  d_value      device array to update
+ * @param[in]   size         number of elements in array
+ * @param[in]   value        value to be added to each element of the buffer
+ *
+ */
+template <typename value_t>
+void transform_increment(raft::handle_t const& handle, value_t* d_value, size_t size, size_t value);
+
 /**
  * @brief    Fill a buffer with a sequence of values
  *
diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index 822a6a2d991..de636f84ff1 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -85,7 +85,7 @@ enum class prior_sources_behavior_t { DEFAULT = 0, CARRY_OVER, EXCLUDE };
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
  * for the sampling.
  * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
  * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
@@ -189,7 +189,7 @@ uniform_neighbor_sample(
  * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_labels Optional device span of labels associted with each starting vertex
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
  * for the sampling.
  * @param label_to_output_comm_rank Optional tuple of device spans mapping label to a particular
  * output rank.  Element 0 of the tuple identifes the label, Element 1 of the tuple identifies the
@@ -312,8 +312,8 @@ struct sampling_flags_t {
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_offsets Optional device span of offsets identifying the range of
- * starting vertex values for this label.
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
+ * for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
@@ -347,7 +347,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,
@@ -400,8 +400,8 @@ heterogeneous_uniform_neighbor_sample(
  * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
- * @param starting_vertex_offsets Optional device span of offsets identifying the range of
- * starting vertex values for this label.
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
+ * for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
@@ -437,7 +437,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,
@@ -485,8 +485,8 @@ heterogeneous_biased_neighbor_sample(
  * @param edge_type_view Optional view object holding edge types for @p graph_view.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
-  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
- * starting vertex values for this label.
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
+ * for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
@@ -519,7 +519,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -570,8 +570,8 @@ homogeneous_uniform_neighbor_sample(
  * corresponding edge can never be selected.
  * @param starting_vertices Device span of starting vertex IDs for the sampling.
  * In a multi-gpu context the starting vertices should be local to this GPU.
-  * @param starting_vertex_offsets Optional device span of offsets identifying the range of
- * starting vertex values for this label.
+ * @param starting_vertex_labels Optional device span of labels associated with each starting vertex
+ * for the sampling.
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
@@ -606,7 +606,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 6827f23e986..7f9ac016d85 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -509,12 +509,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
- * If not specifed the data is not organized in output.  If specified then the all data in the range
- * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_offsets is also specified.
- * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
@@ -539,7 +533,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_sample_heterogeneous_fan_out_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -570,12 +563,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
- * If not specifed the data is not organized in output.  If specified then the all data in the range
- * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_offsets is also specified.
- * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
@@ -600,7 +587,6 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_sample_heterogeneous_fan_out_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -628,12 +614,6 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
- * If not specifed the data is not organized in output.  If specified then the all data in the range
- * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_offsets is also specified.
- * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  sampling_options
@@ -651,7 +631,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -679,12 +658,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
  * @param [in]  start_vertices Device array of start vertices for the sampling
  * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  label_to_comm_rank Device array identifying which comm rank the output for a
- * particular label should be shuffled in the output based on the @p start_vertex_offsets. 
- * If not specifed the data is not organized in output.  If specified then the all data in the range
- * @p start_vertex_offsets[i+1] - @p start_vertex_offsets[i] will be shuffled to rank @p.  This
- * cannot be specified unless @p start_vertex_offsets is also specified.
- * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
  *                            We only support fan_out values of type INT32
  * @param [in]  sampling_options
@@ -703,7 +676,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 04090b22a1d..de3d79875d2 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -83,7 +83,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr}; // start_vertex_label_offsets - local to each GPU - For every start vertex, it needs to pass a global label IDs and label_to_comm_rank (?) - 
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};
@@ -774,7 +774,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   cugraph::c_api::cugraph_edge_property_view_t const* edge_biases_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertices_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_offsets_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   int num_edge_types_{};
   cugraph::c_api::cugraph_sampling_options_t options_{};
@@ -789,7 +788,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
     cugraph_edge_property_view_t const* edge_biases,
     cugraph_type_erased_device_array_view_t const* start_vertices,
     cugraph_type_erased_device_array_view_t const* start_vertex_offsets,
-    cugraph_type_erased_device_array_view_t const* label_to_comm_rank,
     cugraph_type_erased_host_array_view_t const* fan_out,
     int num_edge_types,
     cugraph::c_api::cugraph_sampling_options_t options,
@@ -807,9 +805,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       start_vertex_offsets_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
           start_vertex_offsets)),
-      label_to_comm_rank_(
-        reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
-          label_to_comm_rank)),
       fan_out_(
         reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)),
       num_edge_types_(num_edge_types),
@@ -872,8 +867,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                  handle_.get_stream());
 
       std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
-
       std::optional<rmm::device_uvector<size_t>> start_vertex_offsets{std::nullopt};
+      std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{std::nullopt};
 
       if (start_vertex_offsets_ != nullptr) {
         start_vertex_offsets =
@@ -882,12 +877,49 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                    start_vertex_offsets_->as_type<size_t>(),
                    start_vertex_offsets_->size_,
                    handle_.get_stream());
+        
+
+        // Get the number of labels on each GPU
+        auto num_local_labels = start_vertex_offsets_->size_ - 1;
+        auto global_labels = cugraph::host_scalar_allgather(
+          handle_.get_comms(), num_local_labels, handle_.get_stream());
+        
+        std::exclusive_scan(
+            global_labels.begin(), global_labels.end(), global_labels.begin(), size_t{0});
+        
+        // Compute the global start_vertex_label_offsets
+        cugraph::detail::transform_increment(handle_.get_stream(),
+                                             (*start_vertex_offsets).begin(),
+                                             (*start_vertex_offsets).size(),
+                                             global_labels[handle_.get_comms().get_rank()]);
+
+        // Retrieve the start_vertex_labels
+        start_vertex_labels =
+          cugraph::detail::convert_starting_vertex_offsets_to_labels(
+            handle_,
+            raft::device_span<size_t const>{(*start_vertex_offsets).data(),
+                                            (*start_vertex_offsets).size()});
       }
 
       if constexpr (multi_gpu) {
+        if (start_vertex_labels) {
+
+          rmm::device_uvector<label_t> label_to_comm_rank_d_vector((*start_vertex_labels).size(), handle_.get_stream());
+
+          cugraph::detail::scalar_fill(handle_.get_stream(),
+                                       label_to_comm_rank_d_vector.begin(),
+                                       label_to_comm_rank_d_vector.size(),
+                                       label_t{handle_.get_comms().get_rank()});
+          
+          std::tie(start_vertices, *start_vertex_labels, *label_to_comm_rank) =
+            cugraph::detail::shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices), std::move(*start_vertex_labels), std::move(label_to_comm_rank_d_vector));
+          
+        } else {
           start_vertices =
             cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
               handle_, std::move(start_vertices));
+        }
       }
 
       //
@@ -919,7 +951,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                             vertex_t{0});
       
       // FIXME: For biased sampling, the user should pass either biases or edge weights,
-      // otherwised throw an error and suggest them to use uniform neighbor sample instead
+      // otherwised throw an error and suggest the user to call uniform neighbor sample instead
 
       if (num_edge_types_ > 1) {
         // call heterogeneous neighbor sample
@@ -935,12 +967,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
+                ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
+                                                                      (*start_vertex_labels).size())
                 : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
+              label_to_comm_rank
+                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
+                                                                  (*label_to_comm_rank).size()})
                 : std::nullopt,
               raft::host_span<const int>(
                                         fan_out_->as_type<const int>(), fan_out_->size_),
@@ -963,12 +995,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
+                ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
+                                                                      (*start_vertex_labels).size())
                 : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
+              label_to_comm_rank
+                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
+                                                                  (*label_to_comm_rank).size()})
                 : std::nullopt,
               raft::host_span<const int>(
                                         fan_out_->as_type<const int>(), fan_out_->size_),
@@ -996,12 +1028,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               (edge_biases != nullptr) ? *edge_biases : edge_weights->view(),
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
+                ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
+                                                                   (*start_vertex_labels).size())
                 : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
+              label_to_comm_rank
+                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
+                                                                  (*label_to_comm_rank).size()})
                 : std::nullopt,
               raft::host_span<const int>(
                                         fan_out_->as_type<const int>(), fan_out_->size_),
@@ -1024,12 +1056,12 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<size_t const>>(start_vertex_offsets->data(),
-                                                                      start_vertex_offsets->size())
+                ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
+                                                                      (*start_vertex_labels).size())
                 : std::nullopt,
-              (label_to_comm_rank_ != nullptr)
-                ? std::make_optional(raft::device_span<int const>{label_to_comm_rank_->as_type<int>(),
-                                                    label_to_comm_rank_->size_})
+              label_to_comm_rank
+                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
+                                                                  (*label_to_comm_rank).size()})
                 : std::nullopt,
               raft::host_span<const int>(
                                         fan_out_->as_type<const int>(), fan_out_->size_),
@@ -1129,8 +1161,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
             rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
             rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-            // FIXME: Update this function to handle the new API with 'starting_vertex_offsets'
-            // and indices.
+
             std::tie(majors,
                     output_major_offsets,
                     minors,
@@ -1191,8 +1222,6 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           // extract the edge_type from label_type_hop_offsets
           std::optional<rmm::device_uvector<size_t>> label_type_hop_offsets{std::nullopt};
 
-          // FIXME: Update this function to handle the new API with 'starting_vertex_offsets'
-          // and indices.
           std::tie(output_majors,
               minors,
               wgt,
@@ -1951,7 +1980,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -1974,11 +2002,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
-               *error);
     
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
@@ -2001,7 +2024,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
                                     nullptr,
                                     start_vertices,
                                     start_vertex_offsets,
-                                    label_to_comm_rank,
                                     fan_out,
                                     num_edge_types,
                                     std::move(options_cpp),
@@ -2017,7 +2039,6 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
@@ -2047,11 +2068,6 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
-               *error);
     
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
@@ -2074,7 +2090,6 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_offsets,
-                                    label_to_comm_rank,
                                     fan_out,
                                     num_edge_types,
                                     std::move(options_cpp),
@@ -2089,7 +2104,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -2111,11 +2125,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
-               *error);
   
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
                 fan_out)
@@ -2139,7 +2148,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
                                     nullptr,
                                     start_vertices,
                                     start_vertex_offsets,
-                                    label_to_comm_rank,
                                     fan_out,
                                     1, // num_edge_types
                                     std::move(options_cpp),
@@ -2155,7 +2163,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -2184,11 +2191,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-
-  CAPI_EXPECTS((label_to_comm_rank == nullptr) || (start_vertex_offsets != nullptr),
-               CUGRAPH_INVALID_INPUT,
-               "cannot specify label_to_comm_rank unless start_vertex_offsets is also specified",
-               *error);
   
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
                 fan_out)
@@ -2212,7 +2214,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                                     edge_biases,
                                     start_vertices,
                                     start_vertex_offsets,
-                                    label_to_comm_rank,
                                     fan_out,
                                     1, // num_edge_types
                                     std::move(options_cpp),
diff --git a/cpp/src/detail/utility_wrappers_32.cu b/cpp/src/detail/utility_wrappers_32.cu
index de407f12493..bd5927c8cb3 100644
--- a/cpp/src/detail/utility_wrappers_32.cu
+++ b/cpp/src/detail/utility_wrappers_32.cu
@@ -73,6 +73,16 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             size_t size,
                             uint32_t start_value);
 
+template void transform_increment(rmm::cuda_stream_view const& stream_view,
+                                  int32_t* d_value,
+                                  size_t size,
+                                  size_t value);
+
+template void transform_increment(rmm::cuda_stream_view const& stream_view,
+                                  uint32_t* d_value,
+                                  size_t size,
+                                  size_t value);
+
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int32_t* d_value,
                           size_t size,
diff --git a/cpp/src/detail/utility_wrappers_64.cu b/cpp/src/detail/utility_wrappers_64.cu
index 2c136d5902b..31f021a7ff5 100644
--- a/cpp/src/detail/utility_wrappers_64.cu
+++ b/cpp/src/detail/utility_wrappers_64.cu
@@ -71,6 +71,16 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             size_t size,
                             uint64_t start_value);
 
+template void transform_increment(rmm::cuda_stream_view const& stream_view,
+                                  int64_t* d_value,
+                                  size_t size,
+                                  size_t value);
+
+template void transform_increment(rmm::cuda_stream_view const& stream_view,
+                                  uint64_t* d_value,
+                                  size_t size,
+                                  size_t value);
+
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int64_t* d_value,
                           size_t size,
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index 074d7044261..9f6e5cbb135 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -72,6 +72,22 @@ void sequence_fill(rmm::cuda_stream_view const& stream_view,
   thrust::sequence(rmm::exec_policy(stream_view), d_value, d_value + size, start_value);
 }
 
+template <typename value_t>
+void transform_increment(rmm::cuda_stream_view const& stream_view,
+                         value_t* d_value,
+                         size_t size,
+                         size_t incr)
+{
+  thrust::transform(rmm::exec_policy(stream_view),
+                    d_value,
+                    d_value + size,
+                    d_value,
+                    cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
+                      return static_cast<value_t>(value + incr);
+                    }));
+}
+
+
 template <typename value_t>
 void stride_fill(rmm::cuda_stream_view const& stream_view,
                  value_t* d_value,
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 899ff644179..a5b0ca05a10 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -511,7 +511,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,
@@ -519,11 +519,6 @@ heterogeneous_uniform_neighbor_sample(
   bool do_expensive_check)
 { 
   using bias_t = weight_t; // dummy
-  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
-
-  if (starting_vertex_offsets)
-    starting_vertex_labels =
-      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
   
   auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
     = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
@@ -535,9 +530,7 @@ heterogeneous_uniform_neighbor_sample(
         edge_type_view,
         std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
         starting_vertices,
-        starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
-                                    starting_vertex_labels.data(), starting_vertex_labels.size()})
-                                : std::nullopt,
+        starting_vertex_labels,
         label_to_output_comm_rank,
         fan_out,
         num_edge_types,
@@ -580,18 +573,13 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   edge_type_t num_edge_types,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 { 
-  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
-
-  if (starting_vertex_offsets)
-    starting_vertex_labels =
-      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
   
   auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
     = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
@@ -603,9 +591,7 @@ heterogeneous_biased_neighbor_sample(
         edge_type_view,
         std::make_optional(edge_bias_view),
         starting_vertices,
-        starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
-                                    starting_vertex_labels.data(), starting_vertex_labels.size()})
-                                : std::nullopt,
+        starting_vertex_labels,
         label_to_output_comm_rank,
         fan_out,
         num_edge_types,
@@ -645,21 +631,14 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_t const*>> edge_id_view,
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  // edge_type_t num_edge_types, /*argument not needed for homogeneous neighbor sample*/
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
   using bias_t = weight_t; // dummy
-  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
-
-  if (starting_vertex_offsets)
-    starting_vertex_labels =
-      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
 
-  
   auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
     = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
@@ -670,9 +649,7 @@ homogeneous_uniform_neighbor_sample(
     edge_type_view,
     std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
     starting_vertices,
-    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
-                                starting_vertex_labels.data(), starting_vertex_labels.size()})
-                            : std::nullopt,
+    starting_vertex_labels,
     label_to_output_comm_rank,
     fan_out,
     edge_type_t{1},
@@ -714,21 +691,13 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<edge_t, edge_type_t const*>> edge_type_view,
   edge_property_view_t<edge_t, bias_t const*> edge_bias_view,
   raft::device_span<vertex_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
-  // edge_type_t num_edge_types, /*argument not needed for homogeneous neighbor sample*/
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
 
-  rmm::device_uvector<int32_t> starting_vertex_labels(0, handle.get_stream());
-
-  if (starting_vertex_offsets)
-    starting_vertex_labels =
-      detail::convert_starting_vertex_offsets_to_labels(handle, *starting_vertex_offsets);
-
-  
   auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
     = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
     handle,
@@ -739,9 +708,7 @@ homogeneous_biased_neighbor_sample(
     edge_type_view,
     std::make_optional(edge_bias_view),
     starting_vertices,
-    starting_vertex_offsets ? std::make_optional(raft::device_span<int32_t const>{
-                                starting_vertex_labels.data(), starting_vertex_labels.size()})
-                            : std::nullopt,
+    starting_vertex_labels,
     label_to_output_comm_rank,
     fan_out,
     edge_type_t{1},
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
index 6384f80f8e8..ea2444c1e3c 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
@@ -142,7 +142,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -165,7 +165,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -187,7 +187,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -209,7 +209,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
index 0b153b976c4..846fb33eb21 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
@@ -143,7 +143,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -166,7 +166,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -188,7 +188,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -210,7 +210,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
index da2c89dbeb7..72eec5f6782 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
@@ -142,7 +142,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -165,7 +165,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -187,7 +187,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -209,7 +209,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
index d3f7cf3ce74..46db410f5ae 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
@@ -142,7 +142,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -165,7 +165,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -187,7 +187,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -209,7 +209,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int32_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
index 205749bb9fc..93c0fe176f4 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
@@ -142,7 +142,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -165,7 +165,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -187,7 +187,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -209,7 +209,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
index 97cccb9d307..c3764aa0852 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
@@ -142,7 +142,7 @@ heterogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -165,7 +165,7 @@ heterogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   int32_t num_edge_types,
@@ -187,7 +187,7 @@ homogeneous_uniform_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
@@ -209,7 +209,7 @@ homogeneous_biased_neighbor_sample(
   std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
   edge_property_view_t<int64_t, double const*> edge_bias_view,
   raft::device_span<int64_t const> starting_vertices,
-  std::optional<raft::device_span<size_t const>> starting_vertex_offsets,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
   std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
   raft::host_span<int32_t const> fan_out,
   sampling_flags_t sampling_flags,
diff --git a/cpp/src/utilities/shuffle_vertices.cuh b/cpp/src/utilities/shuffle_vertices.cuh
index adce03f7c29..8707068bb90 100644
--- a/cpp/src/utilities/shuffle_vertices.cuh
+++ b/cpp/src/utilities/shuffle_vertices.cuh
@@ -44,22 +44,41 @@ rmm::device_uvector<vertex_t> shuffle_vertices_by_gpu_id_impl(
   return d_rx_vertices;
 }
 
-template <typename vertex_t, typename value_t, typename func_t>
-std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value_t>>
+template <typename vertex_t, typename value0_t, typename value1_t, typename func_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, std::optional<rmm::device_uvector<value1_t>>>
 shuffle_vertices_and_values_by_gpu_id_impl(raft::handle_t const& handle,
                                            rmm::device_uvector<vertex_t>&& d_vertices,
-                                           rmm::device_uvector<value_t>&& d_values,
+                                           rmm::device_uvector<value0_t>&& d_values_0,
+                                           std::optional<rmm::device_uvector<value1_t>>&& d_values_1,
                                            func_t func)
 {
-  std::tie(d_vertices, d_values, std::ignore) = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
+
+ if (d_values_1) {
+  auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
+    handle.get_comms(),
+    d_vertices.begin(),
+    d_vertices.end(),
+    thrust::make_zip_iterator(d_values_0.begin(), (*d_values_1).begin()),
+    [key_func = func] __device__(auto val) { return key_func(val); },
+    handle.get_stream());
+  
+  return std::make_tuple(
+    std::move(d_shuffled_vertices), std::move(std::get<0>(d_values)), std::make_optional(std::move(std::get<1>(d_values))));
+ } else {
+  auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
     handle.get_comms(),
     d_vertices.begin(),
     d_vertices.end(),
-    d_values.begin(),
+    d_values_0.begin(),
     [key_func = func] __device__(auto val) { return key_func(val); },
     handle.get_stream());
 
-  return std::make_tuple(std::move(d_vertices), std::move(d_values));
+    auto d_values_1 = std::optional<rmm::device_uvector<int32_t>>{std::nullopt};
+
+  return std::make_tuple(
+    std::move(d_shuffled_vertices), std::move(d_values), std::move(d_values_1));
+ }
+  
 }
 
 }  // namespace
@@ -96,12 +115,47 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
   auto const minor_comm_size = minor_comm.get_size();
 
-  return shuffle_vertices_and_values_by_gpu_id_impl(
+  rmm::device_uvector<vertex_t> d_vertices(0, handle.get_stream());
+  rmm::device_uvector<value_t> d_values(0, handle.get_stream());
+
+  std::tie(d_vertices, d_values, std::ignore) = shuffle_vertices_and_values_by_gpu_id_impl(
     handle,
     std::move(vertices),
     std::move(values),
+    std::optional<rmm::device_uvector<int32_t>>{std::nullopt},
+    cugraph::detail::compute_gpu_id_from_ext_vertex_t<vertex_t>{
+      comm_size, major_comm_size, minor_comm_size});
+
+  return std::make_tuple(std::move(d_vertices), std::move(d_values));
+}
+
+template <typename vertex_t, typename value0_t, typename value1_t>
+std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<vertex_t>&& vertices,
+  rmm::device_uvector<value0_t>&& values_0,
+  rmm::device_uvector<value1_t>&& values_1)
+{
+  auto const comm_size       = handle.get_comms().get_size();
+  auto& major_comm           = handle.get_subcomm(cugraph::partition_manager::major_comm_name());
+  auto const major_comm_size = major_comm.get_size();
+  auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
+  auto const minor_comm_size = minor_comm.get_size();
+
+  rmm::device_uvector<vertex_t> d_vertices(0, handle.get_stream());
+  rmm::device_uvector<value0_t> d_values_0(0, handle.get_stream());
+  std::optional<rmm::device_uvector<value1_t>> d_values_1(std::nullopt);
+
+  std::tie(d_vertices, d_values_0, d_values_1) = shuffle_vertices_and_values_by_gpu_id_impl(
+    handle,
+    std::move(vertices),
+    std::move(values_0),
+    std::make_optional(std::move(values_1)),
     cugraph::detail::compute_gpu_id_from_ext_vertex_t<vertex_t>{
       comm_size, major_comm_size, minor_comm_size});
+
+  return std::make_tuple(std::move(d_vertices), std::move(d_values_0), std::move(*d_values_1));
 }
 
 template <typename vertex_t>
@@ -154,17 +208,21 @@ shuffle_int_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
   auto const minor_comm_size = minor_comm.get_size();
 
-  auto return_value = shuffle_vertices_and_values_by_gpu_id_impl(
+  rmm::device_uvector<vertex_t> d_vertices(0, handle.get_stream());
+  rmm::device_uvector<value_t> d_values(0, handle.get_stream());
+
+  std::tie(d_vertices, d_values, std::ignore) = shuffle_vertices_and_values_by_gpu_id_impl(
     handle,
     std::move(vertices),
     std::move(values),
+    std::optional<rmm::device_uvector<int32_t>>{std::nullopt},
     cugraph::detail::compute_gpu_id_from_int_vertex_t<vertex_t>{
       raft::device_span<vertex_t const>(d_vertex_partition_range_lasts.data(),
                                         d_vertex_partition_range_lasts.size()),
       major_comm_size,
       minor_comm_size});
 
-  return return_value;
+  return std::make_tuple(std::move(d_vertices), std::move(d_values));
 }
 
 }  // namespace detail
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
index 1112d7af2f7..f148f71f166 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
@@ -40,6 +40,20 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& vertices,
+  rmm::device_uvector<float>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& vertices,
+  rmm::device_uvector<double>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
index db7be5a3031..c8f158bde6a 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
@@ -62,6 +62,21 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& vertices,
+  rmm::device_uvector<int32_t>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
+template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& vertices,
+  rmm::device_uvector<size_t>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
+
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
index 9592a9c63d1..a61c91b8ac3 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
@@ -40,6 +40,20 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int64_t>&& vertices,
+  rmm::device_uvector<float>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int64_t>&& vertices,
+  rmm::device_uvector<double>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
index 7d968006bc7..1539244307a 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
@@ -54,6 +54,20 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int64_t>&& vertices,
+  rmm::device_uvector<int32_t>&& values_0,
+  rmm::device_uvector<int32_t>&& values_1);
+
+template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
+shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int64_t>&& vertices,
+  rmm::device_uvector<size_t>&& values,
+  rmm::device_uvector<int32_t>&& values_1);
+
 template rmm::device_uvector<int64_t> shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
   raft::handle_t const& handle,
   rmm::device_uvector<int64_t>&& vertices,
diff --git a/python/pylibcugraph/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
index fb07978c70b..90fce23282e 100644
--- a/python/pylibcugraph/pylibcugraph/CMakeLists.txt
+++ b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
@@ -54,8 +54,6 @@ set(cython_sources
     triangle_count.pyx
     two_hop_neighbors.pyx
     uniform_neighbor_sample.pyx
-    #neighbor_sample.pyx // Fix me, break the APi into homogeneous nad heterogeneous neighbor sample
-    #biased_neighbor_sample.pyx
     uniform_random_walks.pyx
     utils.pyx
     weakly_connected_components.pyx

From e357f4291c525b653e09525e5f765b2dba9e36a4 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 25 Sep 2024 14:13:04 -0700
Subject: [PATCH 050/117] remove unsued code

---
 cpp/src/c_api/neighbor_sampling.cpp | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index de3d79875d2..76c35f93b4f 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -942,14 +942,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       std::optional<rmm::device_uvector<int32_t>> hop{std::nullopt};
       std::optional<rmm::device_uvector<label_t>> edge_label{std::nullopt};
       std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
-
-      rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
-
-      cugraph::detail::sequence_fill(handle_.get_stream(),
-                            vertex_type_offsets.begin(),
-                            vertex_type_offsets.size(),
-                            vertex_t{0});
-      
+  
       // FIXME: For biased sampling, the user should pass either biases or edge weights,
       // otherwised throw an error and suggest the user to call uniform neighbor sample instead
 

From 6081978ff687e9018dd04f15d5ccc1465b1b9f12 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 25 Sep 2024 14:46:24 -0700
Subject: [PATCH 051/117] remove outdated comment

---
 cpp/src/c_api/neighbor_sampling.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 76c35f93b4f..ebee588309b 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -83,7 +83,7 @@ struct uniform_neighbor_sampling_functor : public cugraph::c_api::abstract_funct
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* start_vertex_labels_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_list_{nullptr};
   cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_to_comm_rank_{nullptr};
-  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr}; // start_vertex_label_offsets - local to each GPU - For every start vertex, it needs to pass a global label IDs and label_to_comm_rank (?) - 
+  cugraph::c_api::cugraph_type_erased_device_array_view_t const* label_offsets_{nullptr};
   cugraph::c_api::cugraph_type_erased_host_array_view_t const* fan_out_{nullptr};
   cugraph::c_api::cugraph_rng_state_t* rng_state_{nullptr};
   cugraph::c_api::cugraph_sampling_options_t options_{};

From 73b3ffe58d7cd9a75c253c4a8aeaf7e2ae503276 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 25 Sep 2024 15:01:03 -0700
Subject: [PATCH 052/117] remove unnecessary copy

---
 cpp/src/c_api/neighbor_sampling.cpp | 19 +++++--------------
 1 file changed, 5 insertions(+), 14 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index ebee588309b..782ceb9d403 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -867,18 +867,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                  handle_.get_stream());
 
       std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
-      std::optional<rmm::device_uvector<size_t>> start_vertex_offsets{std::nullopt};
       std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{std::nullopt};
 
       if (start_vertex_offsets_ != nullptr) {
-        start_vertex_offsets =
-          rmm::device_uvector<size_t>{start_vertex_offsets_->size_, handle_.get_stream()};
-        raft::copy(start_vertex_offsets->data(),
-                   start_vertex_offsets_->as_type<size_t>(),
-                   start_vertex_offsets_->size_,
-                   handle_.get_stream());
-        
-
         // Get the number of labels on each GPU
         auto num_local_labels = start_vertex_offsets_->size_ - 1;
         auto global_labels = cugraph::host_scalar_allgather(
@@ -889,16 +880,16 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         
         // Compute the global start_vertex_label_offsets
         cugraph::detail::transform_increment(handle_.get_stream(),
-                                             (*start_vertex_offsets).begin(),
-                                             (*start_vertex_offsets).size(),
+                                             start_vertex_offsets_->as_type<size_t>(),
+                                             start_vertex_offsets_->size_,
                                              global_labels[handle_.get_comms().get_rank()]);
 
         // Retrieve the start_vertex_labels
         start_vertex_labels =
           cugraph::detail::convert_starting_vertex_offsets_to_labels(
             handle_,
-            raft::device_span<size_t const>{(*start_vertex_offsets).data(),
-                                            (*start_vertex_offsets).size()});
+            raft::device_span<size_t const>{start_vertex_offsets_->as_type<size_t>(),
+                                            start_vertex_offsets_->size_});
       }
 
       if constexpr (multi_gpu) {
@@ -1246,7 +1237,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               raft::device_span<vertex_t const>{
                     vertex_type_offsets.data(), vertex_type_offsets.size()},
               
-              start_vertex_offsets ? start_vertex_offsets->size() : size_t{1},
+              start_vertex_offsets_ ? start_vertex_offsets_->size_ : size_t{1},
               hop ? fan_out_->size_ : size_t{1},
               size_t{1},
               size_t{1},

From ea972f362615a19f144349fab3f6ba12f2f02161 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 26 Sep 2024 13:40:19 -0700
Subject: [PATCH 053/117] remove outdate arguments

---
 cpp/include/cugraph_c/sampling_algorithms.h | 56 ++-------------------
 cpp/src/c_api/neighbor_sampling.cpp         | 26 ----------
 cpp/src/c_api/neighbor_sampling.hpp         | 27 ----------
 3 files changed, 4 insertions(+), 105 deletions(-)
 delete mode 100644 cpp/src/c_api/neighbor_sampling.hpp

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 7f9ac016d85..3ef176d9080 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -330,51 +330,12 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
 /**
  * @brief     Opaque neighborhood sampling heterogeneous fan_out type
  */
-typedef struct {
-  int32_t align_;
-} cugraph_sample_heterogeneous_fan_out_t;
 
-/**
- * @brief       Create heterogeneous fan_out
- *
- * Input data will be stored in the heterogenous_fan_out.
- * 
- * The fan_out is going to be a CSR structure, the edge_type_offsets will define which range
- * of the fan_out array is associated with each edge type, the fan_out will be the values of
- * fan_out for that hop/type. So for edge type k, fan_out[edge_type_offsets[k]] will identify
- * the fan_out for hop 0 for edge type k.  fan_out[edge_type_offsets[k] +1] will identify the
- * fan_out for hop 1, etc. edge_type_offsets[k+1] will mark the beginning of the fan_out
- * array for type k+1 (and the end of the fan_out array for type k.
- *
- * @param [in]  handle         Handle for accessing resources
- * @param [in]  graph          Pointer to graph
- * @param [in]  edge_type_offsets Type erased array of edge type offsets
- * @param [in]  fan_out         Type erased array of fan_out values
- * @param [out] heterogeneous_fan_out Opaque pointer to fan_out_t
- * @param [out] error          Pointer to an error object storing details of any error.  Will
- *                             be populated if error code is not CUGRAPH_SUCCESS
- * @return error code
- */
-cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
-  const cugraph_resource_handle_t* handle,
-  cugraph_graph_t* graph,
-  const cugraph_type_erased_host_array_view_t* edge_type_offsets,
-  const cugraph_type_erased_host_array_view_t* fan_out,
-  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
-  cugraph_error_t** error);
-
-/**
- * @brief     Free edge type and fan_out pairs
- *
- * @param [in]    heterogeneous_fan_out The edge type size and fan_out values
- */
-void cugraph_heterogeneous_fan_out_free(cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out);
 
 /**
  * @brief     Uniform Neighborhood Sampling
  * 
- * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
- * 'is_biased' set to false instead
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_uniform_neighbor_sample
  *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
@@ -431,8 +392,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
 /**
  * @brief     Biased Neighborhood Sampling
  * 
- * @deprecated  This API will be deleted, use cugraph_homogeneous_neighbor_sample with
- * 'is_biased' set to true instead
+ * @deprecated  This API will be deleted, use cugraph_homogeneous_biased_neighbor_sample.
  * 
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
@@ -516,9 +476,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  *                            sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
- * @param [in]  is_biased
- *                           A flag specifying whether to run biased neighborhood sampling
- *                           (if set to true) or uniform neighbor sampling.
  * @param [in]  do_expensive_check
  *                           A flag to run expensive checks for input arguments (if set to true)
  * @param [out]  result      Output from the uniform_neighbor_sample call
@@ -533,10 +490,9 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_sample_heterogeneous_fan_out_t* fan_out,
+  const cugraph_type_erased_device_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
-  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
@@ -570,9 +526,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  *                            sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
- * @param [in]  is_biased
- *                           A flag specifying whether to run biased neighborhood sampling
- *                           (if set to true) or uniform neighbor sampling.
  * @param [in]  do_expensive_check
  *                           A flag to run expensive checks for input arguments (if set to true)
  * @param [out]  result      Output from the uniform_neighbor_sample call
@@ -587,10 +540,9 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_sample_heterogeneous_fan_out_t* fan_out,
+  const cugraph_type_erased_device_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
-  bool_t is_biased,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 782ceb9d403..39ecec73f96 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-#include "c_api/neighbor_sampling.hpp"
 
 #include "c_api/abstract_functor.hpp"
 #include "c_api/graph.hpp"
@@ -1860,8 +1859,6 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
     CUGRAPH_INVALID_INPUT,
     "fan_out should be of type int",
     *error);
-  
-  bool is_biased = false;
 
   uniform_neighbor_sampling_functor functor{handle,
                                             graph,
@@ -1941,8 +1938,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
-  bool is_biased = true;
-
   biased_neighbor_sampling_functor functor{handle,
                                            graph,
                                            edge_biases,
@@ -2125,7 +2120,6 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
                "vertex type of graph and start_vertices must match",
                *error);
 
-
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
@@ -2205,23 +2199,3 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                                     do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
-
-
-
-
-
-
-extern "C" cugraph_error_code_t cugraph_create_heterogeneous_fan_out(
-  const cugraph_resource_handle_t* handle,
-  cugraph_graph_t* graph,
-  const cugraph_type_erased_host_array_view_t* edge_type_offsets,
-  const cugraph_type_erased_host_array_view_t* fan_out,
-  cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
-  cugraph_error_t** error)
-{
-  *heterogeneous_fan_out = reinterpret_cast<cugraph_sample_heterogeneous_fan_out_t*> (new std::tuple<cugraph::c_api::cugraph_type_erased_host_array_t *, cugraph::c_api::cugraph_type_erased_host_array_t*> {
-          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(edge_type_offsets)),
-          new cugraph::c_api::cugraph_type_erased_host_array_t(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out))});
-  
-  return CUGRAPH_SUCCESS;
-}
diff --git a/cpp/src/c_api/neighbor_sampling.hpp b/cpp/src/c_api/neighbor_sampling.hpp
deleted file mode 100644
index c4a2f8f3575..00000000000
--- a/cpp/src/c_api/neighbor_sampling.hpp
+++ /dev/null
@@ -1,27 +0,0 @@
-/*
- * Copyright (c) 2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include "c_api/array.hpp"
-namespace cugraph {
-namespace c_api {
-
-using cugraph_sample_heterogeneous_fan_out_t = std::tuple<
-  cugraph_type_erased_host_array_t*, cugraph_type_erased_host_array_t*>;
-
-}  // namespace c_api
-}  // namespace cugraph

From 882219254390b6bf63866fe410c50230674aa1a5 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 26 Sep 2024 22:34:34 -0700
Subject: [PATCH 054/117] fix typo

---
 cpp/include/cugraph_c/sampling_algorithms.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 3ef176d9080..9f6e5783bef 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -490,7 +490,7 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* fan_out,
+  const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -540,7 +540,7 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-  const cugraph_type_erased_device_array_view_t* fan_out,
+  const cugraph_type_erased_host_array_view_t* fan_out,
   int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,

From e02a513853381683d4f6f1dd6d60c44265852505 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 26 Sep 2024 22:35:55 -0700
Subject: [PATCH 055/117] update plc API of heterogeneous neighbor sample

---
 .../pylibcugraph/pylibcugraph/CMakeLists.txt  |   1 +
 .../_cugraph_c/sampling_algorithms.pxd        |  64 ++-
 ... heterogeneous_biased_neighbor_sample.pyx} | 191 +++----
 .../heterogeneous_neighbor_sample.pyx         | 511 ------------------
 4 files changed, 99 insertions(+), 668 deletions(-)
 rename python/pylibcugraph/pylibcugraph/{homogeneous_neighbor_sample.pyx => heterogeneous_biased_neighbor_sample.pyx} (71%)
 delete mode 100644 python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx

diff --git a/python/pylibcugraph/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
index 90fce23282e..ddbdc06ad64 100644
--- a/python/pylibcugraph/pylibcugraph/CMakeLists.txt
+++ b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
@@ -63,6 +63,7 @@ set(cython_sources
     all_pairs_sorensen_coefficients.pyx
     all_pairs_overlap_coefficients.pyx
     all_pairs_cosine_coefficients.pyx
+    heterogeneous_biased_neighbor_sample.pyx
 )
 set(linked_libraries cugraph::cugraph;cugraph::cugraph_c)
 
diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index 9366a31605f..ee9dad07e97 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -65,37 +65,59 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_error_t** error
     )
 
-    cdef cugraph_error_code_t cugraph_heterogeneous_neighbor_sample(
+    cdef cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
+        const cugraph_resource_handle_t* handle,
+        cugraph_rng_state_t* rng_state,
+        cugraph_graph_t* graph,
+        const cugraph_type_erased_device_array_view_t* start_vertices,
+        const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+        const cugraph_type_erased_host_array_view_t* fan_out,
+        int num_edge_types,
+        const cugraph_sampling_options_t* options,
+        bool_t do_expensive_check,
+        cugraph_sample_result_t** result,
+        cugraph_error_t** error
+    )
+
+    cdef cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
         const cugraph_resource_handle_t* handle,
         cugraph_rng_state_t* rng_state,
         cugraph_graph_t* graph,
         const cugraph_edge_property_view_t* edge_biases,
         const cugraph_type_erased_device_array_view_t* start_vertices,
-        const cugraph_type_erased_device_array_view_t* start_vertex_labels,
-        const cugraph_type_erased_device_array_view_t* label_list,
-        const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-        const cugraph_type_erased_device_array_view_t* label_offsets,
-        const cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out,
+        const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+        const cugraph_type_erased_host_array_view_t* fan_out,
+        int num_edge_types,
+        const cugraph_sampling_options_t* options,
+        bool_t do_expensive_check,
+        cugraph_sample_result_t** result,
+        cugraph_error_t** error
+    )
+
+    cdef cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
+        const cugraph_resource_handle_t* handle,
+        cugraph_rng_state_t* rng_state,
+        cugraph_graph_t* graph,
+        const cugraph_type_erased_device_array_view_t* start_vertices,
+        const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+        const cugraph_type_erased_device_array_view_t* heterogeneous_fan_out,
+        const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sampling_options_t* options,
-        bool_t is_biased,
         bool_t do_expensive_check,
         cugraph_sample_result_t** result,
         cugraph_error_t** error
     )
 
-    cdef cugraph_error_code_t cugraph_homogeneous_neighbor_sample(
+    cdef cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
         const cugraph_resource_handle_t* handle,
         cugraph_rng_state_t* rng_state,
         cugraph_graph_t* graph,
         const cugraph_edge_property_view_t* edge_biases,
         const cugraph_type_erased_device_array_view_t* start_vertices,
-        const cugraph_type_erased_device_array_view_t* start_vertex_labels,
-        const cugraph_type_erased_device_array_view_t* label_list,
-        const cugraph_type_erased_device_array_view_t* label_to_comm_rank,
-        const cugraph_type_erased_device_array_view_t* label_offsets,
+        const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+        const cugraph_type_erased_device_array_view_t* heterogeneous_fan_out,
         const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sampling_options_t* options,
-        bool_t is_biased,
         bool_t do_expensive_check,
         cugraph_sample_result_t** result,
         cugraph_error_t** error
@@ -142,22 +164,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
             cugraph_error_t** error
         )
 
-    ctypedef struct cugraph_sample_heterogeneous_fan_out_t:
-        pass
-
-    cdef cugraph_error_code_t \
-        cugraph_create_heterogeneous_fan_out(
-            const cugraph_resource_handle_t* handle,
-            cugraph_graph_t* graph,
-            const cugraph_type_erased_host_array_view_t* edge_type_offsets,
-            const cugraph_type_erased_host_array_view_t* fanout,
-            cugraph_sample_heterogeneous_fan_out_t** heterogeneous_fan_out,
-            cugraph_error_t** error
-        )
-
-    cdef void \
-        cugraph_heterogeneous_fan_out_free(
-            cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out);
     # negative sampling
     cdef cugraph_error_code_t \
         cugraph_negative_sampling(
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
similarity index 71%
rename from python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx
rename to python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
index 6b89980f171..6ade76e4edc 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
@@ -55,9 +55,7 @@ from pylibcugraph._cugraph_c.algorithms cimport (
     cugraph_sampling_set_retain_seeds,
 )
 from pylibcugraph._cugraph_c.sampling_algorithms cimport (
-    cugraph_homogeneous_neighbor_sample,
-    cugraph_sample_heterogeneous_fan_out_t,
-    cugraph_heterogeneous_fan_out_free,
+    cugraph_heterogeneous_biased_neighbor_sample,
 )
 from pylibcugraph.resource_handle cimport (
     ResourceHandle,
@@ -83,27 +81,24 @@ from pylibcugraph.random cimport (
 import warnings
 
 # TODO accept cupy/numpy random state in addition to raw seed.
-def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
-                    _GPUGraph input_graph,
-                    start_list,
-                    h_fan_out,
-                    bool_t is_biased,
-                    bool_t with_replacement,
-                    bool_t do_expensive_check,
-                    with_edge_properties=False,
-                    batch_id_list=None,
-                    label_list=None,
-                    label_to_output_comm_rank=None,
-                    label_offsets=None,
-                    prior_sources_behavior=None,
-                    deduplicate_sources=False,
-                    return_hops=False,
-                    renumber=False,
-                    retain_seeds=False,
-                    compression='COO',
-                    compress_per_hop=False,
-                    random_state=None,
-                    return_dict=False,):
+def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
+                                         _GPUGraph input_graph,
+                                         start_vertex_list,
+                                         start_vertex_offsets,
+                                         h_fan_out,
+                                         num_edge_types,
+                                         bool_t with_replacement,
+                                         bool_t do_expensive_check,
+                                         with_edge_properties=False,
+                                         prior_sources_behavior=None,
+                                         deduplicate_sources=False,
+                                         return_hops=False,
+                                         renumber=False,
+                                         retain_seeds=False,
+                                         compression='COO',
+                                         compress_per_hop=False,
+                                         random_state=None,
+                                         return_dict=False,):
     """
     # FIXME: Deprecate uniform_neighbor_sample
     Does both uniform and biased neighborhood sampling, which samples nodes from
@@ -120,20 +115,26 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
 
     input_graph : SGGraph or MGGraph
         The input graph, for either Single or Multi-GPU operations.
+    
+    edge_biases: FIXE: update this - Create edge_biases of type
+        'cugraph_edge_property_view_t' - edge biases not yet supported.
 
-    start_list: device array type
+    start_vertex_list: device array type
         Device array containing the list of starting vertices for sampling.
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
 
-    edge_biases: FIXE: update this - Create edge_biases of type
-    'cugraph_edge_property_view_t' - edge biases not yet supported.
-
-    h_fan_out: numpy array type
+    h_fan_out: tuple of numpy array type
         Device array containing the branching out (fan-out) degrees per
-        starting vertex for each hop level. The sampling method uses the same
-        fan_out value for each type.
+        starting vertex for each hop level
+        
+        The sampling method can use different fan_out values for each edge type
+        which is not the case for homogeneous neighborhood sampling.
     
-    is_biased: bool
-        If false, sampling procedure is done uniform otherwise with biases.
+    num_edge_types: Number of edge types where a value of 1 translates to homogeneous neighbor
+        sample whereas a value greater than 1 translates to heterogeneous neighbor
+        sample.
 
     with_replacement: bool
         If true, sampling procedure is done with replacement (the same vertex
@@ -148,22 +149,6 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
         edges themselves.  Will result in an error if the provided graph
         does not have edge properties.
 
-    batch_id_list: list[int32] (Optional)
-        List of int32 batch ids that is returned with each edge.  Optional
-        argument, defaults to NULL, returning nothing.
-
-    label_list: list[int32] (Optional)
-        List of unique int32 batch ids.  Required if also passing the
-        label_to_output_comm_rank flag.  Default to NULL (does nothing)
-
-    label_to_output_comm_rank: list[int32] (Optional)
-        Maps the unique batch ids in label_list to the rank of the
-        worker that should hold results for that batch id.
-        Defaults to NULL (does nothing)
-
-    label_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
     prior_sources_behavior: str (Optional)
         Options are "carryover", and "exclude".
         Default will leave the source list as-is.
@@ -226,10 +211,7 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
     )
 
     cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
-    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
     cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_size_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
 
     cdef bool_t c_deduplicate_sources = deduplicate_sources
     cdef bool_t c_return_hops = return_hops
@@ -239,17 +221,14 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
     cdef cugraph_error_code_t error_code
     cdef cugraph_error_t* error_ptr
     cdef uintptr_t ai_fan_out_ptr
-    cdef uintptr_t ai_fan_out_size_ptr
-    cdef uintptr_t ai_fan_out_values_ptr
 
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
 
-    assert_CAI_type(start_list, "start_list")
-    assert_CAI_type(batch_id_list, "batch_id_list", True)
-    assert_CAI_type(label_list, "label_list", True)
-    assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
-    assert_CAI_type(label_offsets, "label_offsets", True)
+    assert_CAI_type(start_vertex_list, "start_vertex_list")
+    assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
 
     assert_AI_type(h_fan_out, "h_fan_out")
+
     ai_fan_out_ptr = \
         h_fan_out.__array_interface__["data"][0]
 
@@ -259,76 +238,39 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
             len(h_fan_out),
             get_c_type_from_numpy_type(h_fan_out.dtype))
 
-    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fan_out")
+
 
     cdef cugraph_sample_result_t* result_ptr
 
     cdef uintptr_t cai_start_ptr = \
-        start_list.__cuda_array_interface__["data"][0]
-
-    cdef uintptr_t cai_batch_id_ptr
-    if batch_id_list is not None:
-        cai_batch_id_ptr = \
-            batch_id_list.__cuda_array_interface__['data'][0]
-
-    cdef uintptr_t cai_label_list_ptr
-    if label_list is not None:
-        cai_label_list_ptr = \
-            label_list.__cuda_array_interface__['data'][0]
+        start_vertex_list.__cuda_array_interface__["data"][0]
 
-    cdef uintptr_t cai_label_to_output_comm_rank_ptr
-    if label_to_output_comm_rank is not None:
-        cai_label_to_output_comm_rank_ptr = \
-            label_to_output_comm_rank.__cuda_array_interface__['data'][0]
+    cdef uintptr_t cai_start_vertex_offsets_ptr
+    if start_vertex_offsets is not None:
+        cai_start_vertex_offsets_ptr = \
+            start_vertex_offsets.__cuda_array_interface__['data'][0]
 
-    cdef uintptr_t cai_label_offsets_ptr
-    if label_offsets is not None:
-        cai_label_offsets_ptr = \
-            label_offsets.__cuda_array_interface__['data'][0]
 
-    cdef cugraph_type_erased_device_array_view_t* start_ptr = \
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
         cugraph_type_erased_device_array_view_create(
             <void*>cai_start_ptr,
-            len(start_list),
-            get_c_type_from_numpy_type(start_list.dtype))
-
-    cdef cugraph_type_erased_device_array_view_t* batch_id_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if batch_id_list is not None:
-        batch_id_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_batch_id_ptr,
-                len(batch_id_list),
-                get_c_type_from_numpy_type(batch_id_list.dtype)
-            )
+            len(start_vertex_list),
+            get_c_type_from_numpy_type(start_vertex_list.dtype))
 
-    cdef cugraph_type_erased_device_array_view_t* label_list_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if label_list is not None:
-        label_list_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_list_ptr,
-                len(label_list),
-                get_c_type_from_numpy_type(label_list.dtype)
-            )
 
-    cdef cugraph_type_erased_device_array_view_t* label_to_output_comm_rank_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if label_to_output_comm_rank is not None:
-        label_to_output_comm_rank_ptr = \
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if start_vertex_offsets is not None:
+        start_vertex_offsets_ptr = \
             cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_to_output_comm_rank_ptr,
-                len(label_to_output_comm_rank),
-                get_c_type_from_numpy_type(label_to_output_comm_rank.dtype)
+                <void*>cai_start_vertex_offsets_ptr,
+                len(start_vertex_offsets),
+                get_c_type_from_numpy_type(start_vertex_offsets.dtype)
             )
 
     cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
     if retain_seeds:
-        if label_offsets is None:
+        if start_vertex_offsets is None:
             raise ValueError("Must provide label offsets if retain_seeds is True")
-        label_offsets_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_offsets_ptr,
-                len(label_offsets),
-                get_c_type_from_numpy_type(label_offsets.dtype)
-            )
 
     cg_rng_state = CuGraphRandomState(resource_handle, random_state)
 
@@ -378,37 +320,30 @@ def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
     cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
     cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
 
-    error_code = cugraph_homogeneous_neighbor_sample(
+    error_code = cugraph_heterogeneous_biased_neighbor_sample(
         c_resource_handle_ptr,
         rng_state_ptr,
         c_graph_ptr,
         <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
-        start_ptr,
-        batch_id_ptr,
-        label_list_ptr,
-        label_to_output_comm_rank_ptr,
-        label_offsets_ptr,
+        start_vertex_list_ptr,
+        start_vertex_offsets_ptr,
         fan_out_ptr,
+        num_edge_types,
         sampling_options,
-        is_biased,
         do_expensive_check,
         &result_ptr,
         &error_ptr)
-    assert_success(error_code, error_ptr, "cugraph_homogeneous_neighbor_sample")
+    assert_success(error_code, error_ptr, "cugraph_heterogeneous_biased_neighbor_sample")
 
     # Free the sampling options
     cugraph_sampling_options_free(sampling_options)
 
-    if isinstance(h_fan_out, tuple):
-        cugraph_heterogeneous_fan_out_free(fan_out_ptr)
-
     # Free the two input arrays that are no longer needed.
-    cugraph_type_erased_device_array_view_free(start_ptr)
+    cugraph_type_erased_device_array_view_free(start_vertex_list_ptr)
     cugraph_type_erased_host_array_view_free(fan_out_ptr)
-    if batch_id_list is not None:
-        cugraph_type_erased_device_array_view_free(batch_id_ptr)
-    if label_offsets is not None:
-        cugraph_type_erased_device_array_view_free(label_offsets_ptr)
+
+    if start_vertex_offsets is not None:
+        cugraph_type_erased_device_array_view_free(start_vertex_offsets_ptr)
 
     # Have the SamplingResult instance assume ownership of the result data.
     result = SamplingResult()
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx
deleted file mode 100644
index 0f999adf93d..00000000000
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_neighbor_sample.pyx
+++ /dev/null
@@ -1,511 +0,0 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Have cython use python 3 syntax
-# cython: language_level = 3
-
-from libc.stdint cimport uintptr_t
-
-from pylibcugraph._cugraph_c.resource_handle cimport (
-    bool_t,
-    cugraph_resource_handle_t,
-)
-from pylibcugraph._cugraph_c.properties cimport (
-    cugraph_edge_property_view_t,
-)
-from pylibcugraph._cugraph_c.error cimport (
-    cugraph_error_code_t,
-    cugraph_error_t,
-)
-from pylibcugraph._cugraph_c.array cimport (
-    cugraph_type_erased_device_array_view_t,
-    cugraph_type_erased_device_array_view_create,
-    cugraph_type_erased_device_array_view_free,
-    cugraph_type_erased_host_array_view_t,
-    cugraph_type_erased_host_array_view_create,
-    cugraph_type_erased_host_array_view_free,
-)
-from pylibcugraph._cugraph_c.graph cimport (
-    cugraph_graph_t,
-)
-from pylibcugraph._cugraph_c.algorithms cimport (
-    cugraph_sample_result_t,
-    cugraph_prior_sources_behavior_t,
-    cugraph_compression_type_t,
-    cugraph_sampling_options_t,
-    cugraph_sampling_options_create,
-    cugraph_sampling_options_free,
-    cugraph_sampling_set_with_replacement,
-    cugraph_sampling_set_return_hops,
-    cugraph_sampling_set_prior_sources_behavior,
-    cugraph_sampling_set_dedupe_sources,
-    cugraph_sampling_set_renumber_results,
-    cugraph_sampling_set_compress_per_hop,
-    cugraph_sampling_set_compression_type,
-    cugraph_sampling_set_retain_seeds,
-)
-from pylibcugraph._cugraph_c.sampling_algorithms cimport (
-    cugraph_heterogeneous_neighbor_sample,
-    cugraph_sample_heterogeneous_fan_out_t,
-    cugraph_create_heterogeneous_fan_out,
-    cugraph_heterogeneous_fan_out_free,
-)
-from pylibcugraph.resource_handle cimport (
-    ResourceHandle,
-)
-from pylibcugraph.graphs cimport (
-    _GPUGraph,
-)
-from pylibcugraph.utils cimport (
-    assert_success,
-    assert_CAI_type,
-    assert_AI_type,
-    get_c_type_from_numpy_type,
-)
-from pylibcugraph.internal_types.sampling_result cimport (
-    SamplingResult,
-)
-from pylibcugraph._cugraph_c.random cimport (
-    cugraph_rng_state_t
-)
-from pylibcugraph.random cimport (
-    CuGraphRandomState
-)
-import warnings
-
-# TODO accept cupy/numpy random state in addition to raw seed.
-def heterogeneous_neighbor_sample(ResourceHandle resource_handle,
-                    _GPUGraph input_graph,
-                    start_list,
-                    h_fan_out,
-                    bool_t is_biased,
-                    bool_t with_replacement,
-                    bool_t do_expensive_check,
-                    with_edge_properties=False,
-                    batch_id_list=None,
-                    label_list=None,
-                    label_to_output_comm_rank=None,
-                    label_offsets=None,
-                    prior_sources_behavior=None,
-                    deduplicate_sources=False,
-                    return_hops=False,
-                    renumber=False,
-                    retain_seeds=False,
-                    compression='COO',
-                    compress_per_hop=False,
-                    random_state=None,
-                    return_dict=False,):
-    """
-    # FIXME: Deprecate uniform_neighbor_sample
-    Does both uniform and biased neighborhood sampling, which samples nodes from
-    a graph based on the current node's neighbors, with a corresponding fan_out
-    value at each hop.
-
-    # FIXME: biased neighbor sampling is not yet exposed to PLC.
-
-    Parameters
-    ----------
-    resource_handle: ResourceHandle
-        Handle to the underlying device and host resources needed for
-        referencing data and running algorithms.
-
-    input_graph : SGGraph or MGGraph
-        The input graph, for either Single or Multi-GPU operations.
-
-    start_list: device array type
-        Device array containing the list of starting vertices for sampling.
-
-    edge_biases: FIXE: update this - Create edge_biases of type
-    'cugraph_edge_property_view_t' - edge biases not yet supported.
-
-    h_fan_out: tuple of numpy array type
-        Device array containing the branching out (fan-out) degrees per
-        starting vertex for each hop level in CSR style format. The first
-        element of the tuple is the offset array per edge type id and the second
-        element corresponds to the fan_out values.
-        The sampling method can use different fan_out values for each edge type.
-    
-    is_biased: bool
-        If false, sampling procedure is done uniform otherwise with biases.
-
-    with_replacement: bool
-        If true, sampling procedure is done with replacement (the same vertex
-        can be selected multiple times in the same step).
-
-    do_expensive_check: bool
-        If True, performs more extensive tests on the inputs to ensure
-        validitity, at the expense of increased run time.
-
-    with_edge_properties: bool
-        If True, returns the edge properties of each edges along with the
-        edges themselves.  Will result in an error if the provided graph
-        does not have edge properties.
-
-    batch_id_list: list[int32] (Optional)
-        List of int32 batch ids that is returned with each edge.  Optional
-        argument, defaults to NULL, returning nothing.
-
-    label_list: list[int32] (Optional)
-        List of unique int32 batch ids.  Required if also passing the
-        label_to_output_comm_rank flag.  Default to NULL (does nothing)
-
-    label_to_output_comm_rank: list[int32] (Optional)
-        Maps the unique batch ids in label_list to the rank of the
-        worker that should hold results for that batch id.
-        Defaults to NULL (does nothing)
-
-    label_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
-    prior_sources_behavior: str (Optional)
-        Options are "carryover", and "exclude".
-        Default will leave the source list as-is.
-        Carryover will carry over sources from previous hops to the
-        current hop.
-        Exclude will exclude sources from previous hops from reappearing
-        as sources in future hops.
-
-    deduplicate_sources: bool (Optional)
-        If True, will deduplicate the source list before sampling.
-        Defaults to False.
-
-    renumber: bool (Optional)
-        If True, will renumber the sources and destinations on a
-        per-batch basis and return the renumber map and batch offsets
-        in additional to the standard returns.
-
-    retain_seeds: bool (Optional)
-        If True, will retain the original seeds (original source vertices)
-        in the output even if they do not have outgoing neighbors.
-        Defaults to False.
-
-    compression: str (Optional)
-        Options: COO (default), CSR, CSC, DCSR, DCSR
-        Sets the compression format for the returned samples.
-
-    compress_per_hop: bool (Optional)
-        If False (default), will create a compressed edgelist for the
-        entire batch.
-        If True, will create a separate compressed edgelist per hop within
-        a batch.
-
-    random_state: int (Optional)
-        Random state to use when generating samples.  Optional argument,
-        defaults to a hash of process id, time, and hostname.
-        (See pylibcugraph.random.CuGraphRandomState)
-
-    return_dict: bool (Optional)
-        Whether to return a dictionary instead of a tuple.
-        Optional argument, defaults to False, returning a tuple.
-        This argument will eventually be deprecated in favor
-        of always returning a dictionary.
-
-    Returns
-    -------
-    A tuple of device arrays, where the first and second items in the tuple
-    are device arrays containing the starting and ending vertices of each
-    walk respectively, the third item in the tuple is a device array
-    containing the start labels, and the fourth item in the tuple is a device
-    array containing the indices for reconstructing paths.
-
-    If renumber was set to True, then the fifth item in the tuple is a device
-    array containing the renumber map, and the sixth item in the tuple is a
-    device array containing the renumber map offsets (which delineate where
-    the renumber map for each batch starts).
-
-    """
-    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
-        resource_handle.c_resource_handle_ptr
-    )
-
-    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
-    cdef cugraph_sample_heterogeneous_fan_out_t* heterogeneous_fan_out_ptr = <cugraph_sample_heterogeneous_fan_out_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_offsets_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-    cdef cugraph_type_erased_host_array_view_t* fan_out_values_ptr = <cugraph_type_erased_host_array_view_t*>NULL
-
-    cdef bool_t c_deduplicate_sources = deduplicate_sources
-    cdef bool_t c_return_hops = return_hops
-    cdef bool_t c_renumber = renumber
-    cdef bool_t c_compress_per_hop = compress_per_hop
-
-    cdef cugraph_error_code_t error_code
-    cdef cugraph_error_t* error_ptr
-    cdef uintptr_t ai_fan_out_ptr
-    cdef uintptr_t ai_fan_out_offsets_ptr
-    cdef uintptr_t ai_fan_out_values_ptr
-
-
-    assert_CAI_type(start_list, "start_list")
-    assert_CAI_type(batch_id_list, "batch_id_list", True)
-    assert_CAI_type(label_list, "label_list", True)
-    assert_CAI_type(label_to_output_comm_rank, "label_to_output_comm_rank", True)
-    assert_CAI_type(label_offsets, "label_offsets", True)
-
-    assert_AI_type(h_fan_out[0], "h_fan_out_size")
-    assert_AI_type(h_fan_out[1], "h_fan_out_values")
-    ai_fan_out_offsets_ptr = \
-        h_fan_out[0].__array_interface__["data"][0]
-    ai_fan_out_values_ptr = \
-        h_fan_out[1].__array_interface__["data"][0]
-
-    fan_out_offsets_ptr = \
-        cugraph_type_erased_host_array_view_create(
-            <void*>ai_fan_out_offsets_ptr,
-            len(h_fan_out[0]),
-            get_c_type_from_numpy_type(h_fan_out[0].dtype))
-
-    fan_out_values_ptr = \
-        cugraph_type_erased_host_array_view_create(
-            <void*>ai_fan_out_values_ptr,
-            len(h_fan_out[1]),
-            get_c_type_from_numpy_type(h_fan_out[1].dtype))
-
-    error_code = cugraph_create_heterogeneous_fan_out(
-        c_resource_handle_ptr,
-        c_graph_ptr,
-        fan_out_offsets_ptr,
-        fan_out_values_ptr,
-        &heterogeneous_fan_out_ptr,
-        &error_ptr
-    )
-
-    assert_success(error_code, error_ptr, "cugraph_create_heterogeneous_fan_out")
-
-    cdef cugraph_sample_result_t* result_ptr
-
-    cdef uintptr_t cai_start_ptr = \
-        start_list.__cuda_array_interface__["data"][0]
-
-    cdef uintptr_t cai_batch_id_ptr
-    if batch_id_list is not None:
-        cai_batch_id_ptr = \
-            batch_id_list.__cuda_array_interface__['data'][0]
-
-    cdef uintptr_t cai_label_list_ptr
-    if label_list is not None:
-        cai_label_list_ptr = \
-            label_list.__cuda_array_interface__['data'][0]
-
-    cdef uintptr_t cai_label_to_output_comm_rank_ptr
-    if label_to_output_comm_rank is not None:
-        cai_label_to_output_comm_rank_ptr = \
-            label_to_output_comm_rank.__cuda_array_interface__['data'][0]
-
-    cdef uintptr_t cai_label_offsets_ptr
-    if label_offsets is not None:
-        cai_label_offsets_ptr = \
-            label_offsets.__cuda_array_interface__['data'][0]
-
-    cdef cugraph_type_erased_device_array_view_t* start_ptr = \
-        cugraph_type_erased_device_array_view_create(
-            <void*>cai_start_ptr,
-            len(start_list),
-            get_c_type_from_numpy_type(start_list.dtype))
-
-    cdef cugraph_type_erased_device_array_view_t* batch_id_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if batch_id_list is not None:
-        batch_id_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_batch_id_ptr,
-                len(batch_id_list),
-                get_c_type_from_numpy_type(batch_id_list.dtype)
-            )
-
-    cdef cugraph_type_erased_device_array_view_t* label_list_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if label_list is not None:
-        label_list_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_list_ptr,
-                len(label_list),
-                get_c_type_from_numpy_type(label_list.dtype)
-            )
-
-    cdef cugraph_type_erased_device_array_view_t* label_to_output_comm_rank_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if label_to_output_comm_rank is not None:
-        label_to_output_comm_rank_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_to_output_comm_rank_ptr,
-                len(label_to_output_comm_rank),
-                get_c_type_from_numpy_type(label_to_output_comm_rank.dtype)
-            )
-
-    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
-    if retain_seeds:
-        if label_offsets is None:
-            raise ValueError("Must provide label offsets if retain_seeds is True")
-        label_offsets_ptr = \
-            cugraph_type_erased_device_array_view_create(
-                <void*>cai_label_offsets_ptr,
-                len(label_offsets),
-                get_c_type_from_numpy_type(label_offsets.dtype)
-            )
-
-    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
-
-    cdef cugraph_rng_state_t* rng_state_ptr = \
-        cg_rng_state.rng_state_ptr
-
-    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
-    if prior_sources_behavior is None:
-        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
-    elif prior_sources_behavior == 'carryover':
-        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
-    elif prior_sources_behavior == 'exclude':
-        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
-    else:
-        raise ValueError(
-            f'Invalid option {prior_sources_behavior}'
-            ' for prior sources behavior'
-        )
-
-    cdef cugraph_compression_type_t compression_behavior_e
-    if compression is None or compression == 'COO':
-        compression_behavior_e = cugraph_compression_type_t.COO
-    elif compression == 'CSR':
-        compression_behavior_e = cugraph_compression_type_t.CSR
-    elif compression == 'CSC':
-        compression_behavior_e = cugraph_compression_type_t.CSC
-    elif compression == 'DCSR':
-        compression_behavior_e = cugraph_compression_type_t.DCSR
-    elif compression == 'DCSC':
-        compression_behavior_e = cugraph_compression_type_t.DCSC
-    else:
-        raise ValueError(
-            f'Invalid option {compression}'
-            ' for compression type'
-        )
-
-    cdef cugraph_sampling_options_t* sampling_options
-    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
-    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
-
-    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
-    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
-    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
-    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
-    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
-    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
-    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
-    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
-
-    error_code = cugraph_heterogeneous_neighbor_sample(
-        c_resource_handle_ptr,
-        rng_state_ptr,
-        c_graph_ptr,
-        <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
-        start_ptr,
-        batch_id_ptr,
-        label_list_ptr,
-        label_to_output_comm_rank_ptr,
-        label_offsets_ptr,
-        heterogeneous_fan_out_ptr,
-        sampling_options,
-        is_biased,
-        do_expensive_check,
-        &result_ptr,
-        &error_ptr)
-    assert_success(error_code, error_ptr, "cugraph_heterogeneous_neighbor_sample")
-
-    # Free the sampling options
-    cugraph_sampling_options_free(sampling_options)
-
-    if isinstance(h_fan_out, tuple):
-        cugraph_heterogeneous_fan_out_free(heterogeneous_fan_out_ptr)
-
-    # Free the two input arrays that are no longer needed.
-    cugraph_type_erased_device_array_view_free(start_ptr)
-    cugraph_type_erased_host_array_view_free(fan_out_ptr)
-    if batch_id_list is not None:
-        cugraph_type_erased_device_array_view_free(batch_id_ptr)
-    if label_offsets is not None:
-        cugraph_type_erased_device_array_view_free(label_offsets_ptr)
-
-    # Have the SamplingResult instance assume ownership of the result data.
-    result = SamplingResult()
-    result.set_ptr(result_ptr)
-
-    # Get cupy "views" of the individual arrays to return. These each increment
-    # the refcount on the SamplingResult instance which will keep the data alive
-    # until all references are removed and the GC runs.
-    # TODO Return everything that isn't null in release 23.12
-    if with_edge_properties:
-        cupy_majors = result.get_majors()
-        cupy_major_offsets = result.get_major_offsets()
-        cupy_minors = result.get_minors()
-        cupy_edge_weights = result.get_edge_weights()
-        cupy_edge_ids = result.get_edge_ids()
-        cupy_edge_types = result.get_edge_types()
-        cupy_batch_ids = result.get_batch_ids()
-        cupy_label_hop_offsets = result.get_label_hop_offsets()
-
-        if renumber:
-            cupy_renumber_map = result.get_renumber_map()
-            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
-            # TODO drop the placeholder for hop ids in release 23.12
-            if return_dict:
-                return {
-                    'major_offsets': cupy_major_offsets,
-                    'majors': cupy_majors,
-                    'minors': cupy_minors,
-                    'weight': cupy_edge_weights,
-                    'edge_id': cupy_edge_ids,
-                    'edge_type': cupy_edge_types,
-                    'batch_id': cupy_batch_ids,
-                    'label_hop_offsets': cupy_label_hop_offsets,
-                    'hop_id': None,
-                    'renumber_map': cupy_renumber_map,
-                    'renumber_map_offsets': cupy_renumber_map_offsets
-                }
-            else:
-                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
-                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
-        else:
-            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
-            if return_dict:
-                return {
-                    'major_offsets': cupy_major_offsets,
-                    'majors': cupy_majors,
-                    'minors': cupy_minors,
-                    'weight': cupy_edge_weights,
-                    'edge_id': cupy_edge_ids,
-                    'edge_type': cupy_edge_types,
-                    'batch_id': cupy_batch_ids,
-                    'label_hop_offsets': cupy_label_hop_offsets,
-                    'hop_id': cupy_hop_ids,
-                }
-            else:
-                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
-                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
-
-    else:
-        # TODO this is deprecated, remove it in release 23.12
-        warnings.warn(
-            "Calling uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
-            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
-            "with_edge_properties=True option, returning whatever properties are in the graph.",
-            FutureWarning,
-        )
-
-        cupy_sources = result.get_sources()
-        cupy_destinations = result.get_destinations()
-        cupy_indices = result.get_indices()
-
-        if return_dict:
-            return {
-                'sources': cupy_sources,
-                'destinations': cupy_destinations,
-                'indices': cupy_indices
-            }
-        else:
-            return (cupy_sources, cupy_destinations, cupy_indices)

From d6cb1d5286aec39fd9fb04c426e3ccb6c95afe22 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 09:20:56 -0700
Subject: [PATCH 056/117] fix typo

---
 cpp/include/cugraph_c/sampling_algorithms.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index 9f6e5783bef..abc8f1c4082 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -487,7 +487,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
-  const cugraph_edge_property_view_t* edge_biases,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,

From 54fa15581bf5de33cc0578788e33b24725c88e9c Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 09:21:54 -0700
Subject: [PATCH 057/117] change back the fanout type from a sparse to a dense
 structure

---
 .../sampling/uniform_neighbor_sample.py       | 21 +++----------------
 1 file changed, 3 insertions(+), 18 deletions(-)

diff --git a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
index 5083d50e7fb..2e3d032a20f 100644
--- a/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/uniform_neighbor_sample.py
@@ -12,7 +12,6 @@
 # limitations under the License.
 
 from __future__ import annotations
-from itertools import chain
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import uniform_neighbor_sample as pylibcugraph_uniform_neighbor_sample
@@ -91,10 +90,9 @@ def uniform_neighbor_sample(
     start_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
 
-    fanout_vals : list (int32) or dict
+    fanout_vals : list (int32)
         List of branching out (fan-out) degrees per starting vertex for each
-        hop level or dictionary of edge type and fanout values for
-        heterogeneous fanout type.
+        hop level.
 
     with_replacement: bool, optional (default=True)
         Flag to specify if the random sampling is done with replacement
@@ -316,21 +314,8 @@ def uniform_neighbor_sample(
         fanout_vals = fanout_vals.get().astype("int32")
     elif isinstance(fanout_vals, cudf.Series):
         fanout_vals = fanout_vals.values_host.astype("int32")
-    elif isinstance(fanout_vals, dict):
-        # FIXME: Add expensive check to ensure all dict values are lists
-        # Convert to a tuple of sequence (edge type size and fanout values)
-        edge_type_offsets = []
-        [edge_type_offsets.append(len(s)) for s in list(fanout_vals.values())]
-        edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
-        fanout_vals = (
-            numpy.asarray(edge_type_offsets, dtype="int32"),
-            numpy.asarray(edge_type_fanout_vals, dtype="int32"),
-        )
     else:
-        raise TypeError(
-            "fanout_vals must be a sequence or a dictionary, "
-            f"got: {type(fanout_vals)}"
-        )
+        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
 
     if "weights" in G.edgelist.edgelist_df:
         weight_t = G.edgelist.edgelist_df["weights"].dtype

From 499e041ff05174bbe1ad0ea59152d9e0c9f9073d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 09:22:50 -0700
Subject: [PATCH 058/117] fix typo

---
 .../pylibcugraph/_cugraph_c/sampling_algorithms.pxd             | 2 --
 1 file changed, 2 deletions(-)

diff --git a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
index ee9dad07e97..2af8e1d6d7e 100644
--- a/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
+++ b/python/pylibcugraph/pylibcugraph/_cugraph_c/sampling_algorithms.pxd
@@ -100,7 +100,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         cugraph_graph_t* graph,
         const cugraph_type_erased_device_array_view_t* start_vertices,
         const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-        const cugraph_type_erased_device_array_view_t* heterogeneous_fan_out,
         const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sampling_options_t* options,
         bool_t do_expensive_check,
@@ -115,7 +114,6 @@ cdef extern from "cugraph_c/sampling_algorithms.h":
         const cugraph_edge_property_view_t* edge_biases,
         const cugraph_type_erased_device_array_view_t* start_vertices,
         const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
-        const cugraph_type_erased_device_array_view_t* heterogeneous_fan_out,
         const cugraph_type_erased_host_array_view_t* fan_out,
         const cugraph_sampling_options_t* options,
         bool_t do_expensive_check,

From b571deb3cff5d8563403d65f5838ab509a7343e9 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 09:24:07 -0700
Subject: [PATCH 059/117] add implementation of heterogeneous/homogeneous
 biased/uniform neighborhood sampling

---
 .../pylibcugraph/pylibcugraph/CMakeLists.txt  |   3 +
 .../heterogeneous_biased_neighbor_sample.pyx  |  10 +-
 .../heterogeneous_uniform_neighbor_sample.pyx | 416 +++++++++++++++++
 .../homogeneous_biased_neighbor_sample.pyx    | 418 ++++++++++++++++++
 .../homogeneous_uniform_neighbor_sample.pyx   | 411 +++++++++++++++++
 5 files changed, 1252 insertions(+), 6 deletions(-)
 create mode 100644 python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
 create mode 100644 python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
 create mode 100644 python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx

diff --git a/python/pylibcugraph/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
index ddbdc06ad64..21446d639fe 100644
--- a/python/pylibcugraph/pylibcugraph/CMakeLists.txt
+++ b/python/pylibcugraph/pylibcugraph/CMakeLists.txt
@@ -64,6 +64,9 @@ set(cython_sources
     all_pairs_overlap_coefficients.pyx
     all_pairs_cosine_coefficients.pyx
     heterogeneous_biased_neighbor_sample.pyx
+    heterogeneous_uniform_neighbor_sample.pyx
+    homogeneous_biased_neighbor_sample.pyx
+    homogeneous_uniform_neighbor_sample.pyx
 )
 set(linked_libraries cugraph::cugraph;cugraph::cugraph_c)
 
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
index 6ade76e4edc..53de709e28d 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
@@ -100,12 +100,10 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
                                          random_state=None,
                                          return_dict=False,):
     """
-    # FIXME: Deprecate uniform_neighbor_sample
-    Does both uniform and biased neighborhood sampling, which samples nodes from
+    # FIXME: Deprecate biased_neighbor_sample
+    Performs biased neighborhood sampling, which samples nodes from
     a graph based on the current node's neighbors, with a corresponding fan_out
-    value at each hop.
-
-    # FIXME: biased neighbor sampling is not yet exposed to PLC.
+    value at each hop. The edges are sampled with biases.
 
     Parameters
     ----------
@@ -405,7 +403,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
     else:
         # TODO this is deprecated, remove it in release 23.12
         warnings.warn(
-            "Calling uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            "Calling heterogeneous_biased_neighbor_sample with the 'with_edge_properties' argument is deprecated."
             " Starting in release 23.12, this argument will be removed in favor of behaving like the "
             "with_edge_properties=True option, returning whatever properties are in the graph.",
             FutureWarning,
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
new file mode 100644
index 00000000000..2bcf0d856fc
--- /dev/null
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
@@ -0,0 +1,416 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    bool_t,
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.algorithms cimport (
+    cugraph_sample_result_t,
+    cugraph_prior_sources_behavior_t,
+    cugraph_compression_type_t,
+    cugraph_sampling_options_t,
+    cugraph_sampling_options_create,
+    cugraph_sampling_options_free,
+    cugraph_sampling_set_with_replacement,
+    cugraph_sampling_set_return_hops,
+    cugraph_sampling_set_prior_sources_behavior,
+    cugraph_sampling_set_dedupe_sources,
+    cugraph_sampling_set_renumber_results,
+    cugraph_sampling_set_compress_per_hop,
+    cugraph_sampling_set_compression_type,
+    cugraph_sampling_set_retain_seeds,
+)
+from pylibcugraph._cugraph_c.sampling_algorithms cimport (
+    cugraph_heterogeneous_uniform_neighbor_sample,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+)
+from pylibcugraph.internal_types.sampling_result cimport (
+    SamplingResult,
+)
+from pylibcugraph._cugraph_c.random cimport (
+    cugraph_rng_state_t
+)
+from pylibcugraph.random cimport (
+    CuGraphRandomState
+)
+import warnings
+
+# TODO accept cupy/numpy random state in addition to raw seed.
+def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
+                                          _GPUGraph input_graph,
+                                          start_vertex_list,
+                                          start_vertex_offsets,
+                                          h_fan_out,
+                                          num_edge_types,
+                                          bool_t with_replacement,
+                                          bool_t do_expensive_check,
+                                          with_edge_properties=False,
+                                          prior_sources_behavior=None,
+                                          deduplicate_sources=False,
+                                          return_hops=False,
+                                          renumber=False,
+                                          retain_seeds=False,
+                                          compression='COO',
+                                          compress_per_hop=False,
+                                          random_state=None,
+                                          return_dict=False,):
+    """
+    # FIXME: Deprecate uniform_neighbor_sample
+    Performs biased neighborhood sampling, which samples nodes from
+    a graph based on the current node's neighbors, with a corresponding fan_out
+    value at each hop. The edges are sampled uniformly.
+
+    Parameters
+    ----------
+    resource_handle: ResourceHandle
+        Handle to the underlying device and host resources needed for
+        referencing data and running algorithms.
+
+    input_graph : SGGraph or MGGraph
+        The input graph, for either Single or Multi-GPU operations.
+
+    start_vertex_list: device array type
+        Device array containing the list of starting vertices for sampling.
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    h_fan_out: tuple of numpy array type
+        Device array containing the branching out (fan-out) degrees per
+        starting vertex for each hop level
+        
+        The sampling method can use different fan_out values for each edge type
+        which is not the case for homogeneous neighborhood sampling.
+    
+    num_edge_types: Number of edge types where a value of 1 translates to homogeneous neighbor
+        sample whereas a value greater than 1 translates to heterogeneous neighbor
+        sample.
+
+    with_replacement: bool
+        If true, sampling procedure is done with replacement (the same vertex
+        can be selected multiple times in the same step).
+
+    do_expensive_check: bool
+        If True, performs more extensive tests on the inputs to ensure
+        validitity, at the expense of increased run time.
+
+    with_edge_properties: bool
+        If True, returns the edge properties of each edges along with the
+        edges themselves.  Will result in an error if the provided graph
+        does not have edge properties.
+
+    prior_sources_behavior: str (Optional)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+
+    deduplicate_sources: bool (Optional)
+        If True, will deduplicate the source list before sampling.
+        Defaults to False.
+
+    renumber: bool (Optional)
+        If True, will renumber the sources and destinations on a
+        per-batch basis and return the renumber map and batch offsets
+        in additional to the standard returns.
+
+    retain_seeds: bool (Optional)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+        Defaults to False.
+
+    compression: str (Optional)
+        Options: COO (default), CSR, CSC, DCSR, DCSR
+        Sets the compression format for the returned samples.
+
+    compress_per_hop: bool (Optional)
+        If False (default), will create a compressed edgelist for the
+        entire batch.
+        If True, will create a separate compressed edgelist per hop within
+        a batch.
+
+    random_state: int (Optional)
+        Random state to use when generating samples.  Optional argument,
+        defaults to a hash of process id, time, and hostname.
+        (See pylibcugraph.random.CuGraphRandomState)
+
+    return_dict: bool (Optional)
+        Whether to return a dictionary instead of a tuple.
+        Optional argument, defaults to False, returning a tuple.
+        This argument will eventually be deprecated in favor
+        of always returning a dictionary.
+
+    Returns
+    -------
+    A tuple of device arrays, where the first and second items in the tuple
+    are device arrays containing the starting and ending vertices of each
+    walk respectively, the third item in the tuple is a device array
+    containing the start labels, and the fourth item in the tuple is a device
+    array containing the indices for reconstructing paths.
+
+    If renumber was set to True, then the fifth item in the tuple is a device
+    array containing the renumber map, and the sixth item in the tuple is a
+    device array containing the renumber map offsets (which delineate where
+    the renumber map for each batch starts).
+
+    """
+    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
+        resource_handle.c_resource_handle_ptr
+    )
+
+    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+
+    cdef bool_t c_deduplicate_sources = deduplicate_sources
+    cdef bool_t c_return_hops = return_hops
+    cdef bool_t c_renumber = renumber
+    cdef bool_t c_compress_per_hop = compress_per_hop
+
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+
+    assert_CAI_type(start_vertex_list, "start_vertex_list")
+    assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
+
+    assert_AI_type(h_fan_out, "h_fan_out")
+
+    ai_fan_out_ptr = \
+        h_fan_out.__array_interface__["data"][0]
+
+    fan_out_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_ptr,
+            len(h_fan_out),
+            get_c_type_from_numpy_type(h_fan_out.dtype))
+
+
+
+    cdef cugraph_sample_result_t* result_ptr
+
+    cdef uintptr_t cai_start_ptr = \
+        start_vertex_list.__cuda_array_interface__["data"][0]
+
+    cdef uintptr_t cai_start_vertex_offsets_ptr
+    if start_vertex_offsets is not None:
+        cai_start_vertex_offsets_ptr = \
+            start_vertex_offsets.__cuda_array_interface__['data'][0]
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
+        cugraph_type_erased_device_array_view_create(
+            <void*>cai_start_ptr,
+            len(start_vertex_list),
+            get_c_type_from_numpy_type(start_vertex_list.dtype))
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if start_vertex_offsets is not None:
+        start_vertex_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_start_vertex_offsets_ptr,
+                len(start_vertex_offsets),
+                get_c_type_from_numpy_type(start_vertex_offsets.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if retain_seeds:
+        if start_vertex_offsets is None:
+            raise ValueError("Must provide label offsets if retain_seeds is True")
+
+    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
+
+    cdef cugraph_rng_state_t* rng_state_ptr = \
+        cg_rng_state.rng_state_ptr
+
+    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
+    if prior_sources_behavior is None:
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
+    elif prior_sources_behavior == 'carryover':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
+    elif prior_sources_behavior == 'exclude':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
+    else:
+        raise ValueError(
+            f'Invalid option {prior_sources_behavior}'
+            ' for prior sources behavior'
+        )
+
+    cdef cugraph_compression_type_t compression_behavior_e
+    if compression is None or compression == 'COO':
+        compression_behavior_e = cugraph_compression_type_t.COO
+    elif compression == 'CSR':
+        compression_behavior_e = cugraph_compression_type_t.CSR
+    elif compression == 'CSC':
+        compression_behavior_e = cugraph_compression_type_t.CSC
+    elif compression == 'DCSR':
+        compression_behavior_e = cugraph_compression_type_t.DCSR
+    elif compression == 'DCSC':
+        compression_behavior_e = cugraph_compression_type_t.DCSC
+    else:
+        raise ValueError(
+            f'Invalid option {compression}'
+            ' for compression type'
+        )
+
+    cdef cugraph_sampling_options_t* sampling_options
+    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
+
+    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
+    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
+    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
+    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
+    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
+    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
+    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
+    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
+
+    error_code = cugraph_heterogeneous_uniform_neighbor_sample(
+        c_resource_handle_ptr,
+        rng_state_ptr,
+        c_graph_ptr,
+        start_vertex_list_ptr,
+        start_vertex_offsets_ptr,
+        fan_out_ptr,
+        num_edge_types,
+        sampling_options,
+        do_expensive_check,
+        &result_ptr,
+        &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_heterogeneous_uniform_neighbor_sample")
+
+    # Free the sampling options
+    cugraph_sampling_options_free(sampling_options)
+
+    # Free the two input arrays that are no longer needed.
+    cugraph_type_erased_device_array_view_free(start_vertex_list_ptr)
+    cugraph_type_erased_host_array_view_free(fan_out_ptr)
+
+    if start_vertex_offsets is not None:
+        cugraph_type_erased_device_array_view_free(start_vertex_offsets_ptr)
+
+    # Have the SamplingResult instance assume ownership of the result data.
+    result = SamplingResult()
+    result.set_ptr(result_ptr)
+
+    # Get cupy "views" of the individual arrays to return. These each increment
+    # the refcount on the SamplingResult instance which will keep the data alive
+    # until all references are removed and the GC runs.
+    # TODO Return everything that isn't null in release 23.12
+    if with_edge_properties:
+        cupy_majors = result.get_majors()
+        cupy_major_offsets = result.get_major_offsets()
+        cupy_minors = result.get_minors()
+        cupy_edge_weights = result.get_edge_weights()
+        cupy_edge_ids = result.get_edge_ids()
+        cupy_edge_types = result.get_edge_types()
+        cupy_batch_ids = result.get_batch_ids()
+        cupy_label_hop_offsets = result.get_label_hop_offsets()
+
+        if renumber:
+            cupy_renumber_map = result.get_renumber_map()
+            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+            # TODO drop the placeholder for hop ids in release 23.12
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': None,
+                    'renumber_map': cupy_renumber_map,
+                    'renumber_map_offsets': cupy_renumber_map_offsets
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+        else:
+            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': cupy_hop_ids,
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
+
+    else:
+        # TODO this is deprecated, remove it in release 23.12
+        warnings.warn(
+            "Calling heterogeneous_uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
+            "with_edge_properties=True option, returning whatever properties are in the graph.",
+            FutureWarning,
+        )
+
+        cupy_sources = result.get_sources()
+        cupy_destinations = result.get_destinations()
+        cupy_indices = result.get_indices()
+
+        if return_dict:
+            return {
+                'sources': cupy_sources,
+                'destinations': cupy_destinations,
+                'indices': cupy_indices
+            }
+        else:
+            return (cupy_sources, cupy_destinations, cupy_indices)
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
new file mode 100644
index 00000000000..926d907995e
--- /dev/null
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
@@ -0,0 +1,418 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    bool_t,
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.properties cimport (
+    cugraph_edge_property_view_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.algorithms cimport (
+    cugraph_sample_result_t,
+    cugraph_prior_sources_behavior_t,
+    cugraph_compression_type_t,
+    cugraph_sampling_options_t,
+    cugraph_sampling_options_create,
+    cugraph_sampling_options_free,
+    cugraph_sampling_set_with_replacement,
+    cugraph_sampling_set_return_hops,
+    cugraph_sampling_set_prior_sources_behavior,
+    cugraph_sampling_set_dedupe_sources,
+    cugraph_sampling_set_renumber_results,
+    cugraph_sampling_set_compress_per_hop,
+    cugraph_sampling_set_compression_type,
+    cugraph_sampling_set_retain_seeds,
+)
+from pylibcugraph._cugraph_c.sampling_algorithms cimport (
+    cugraph_homogeneous_biased_neighbor_sample,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+)
+from pylibcugraph.internal_types.sampling_result cimport (
+    SamplingResult,
+)
+from pylibcugraph._cugraph_c.random cimport (
+    cugraph_rng_state_t
+)
+from pylibcugraph.random cimport (
+    CuGraphRandomState
+)
+import warnings
+
+# TODO accept cupy/numpy random state in addition to raw seed.
+def homogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
+                                       _GPUGraph input_graph,
+                                       start_vertex_list,
+                                       start_vertex_offsets,
+                                       h_fan_out,
+                                       num_edge_types,
+                                       bool_t with_replacement,
+                                       bool_t do_expensive_check,
+                                       with_edge_properties=False,
+                                       prior_sources_behavior=None,
+                                       deduplicate_sources=False,
+                                       return_hops=False,
+                                       renumber=False,
+                                       retain_seeds=False,
+                                       compression='COO',
+                                       compress_per_hop=False,
+                                       random_state=None,
+                                       return_dict=False):
+    """
+    # FIXME: Deprecate biased_neighbor_sample
+    Performs biased neighborhood sampling, which samples nodes from
+    a graph based on the current node's neighbors, with a corresponding fan_out
+    value at each hop. The edges are sampled with biases.
+
+    Parameters
+    ----------
+    resource_handle: ResourceHandle
+        Handle to the underlying device and host resources needed for
+        referencing data and running algorithms.
+
+    input_graph : SGGraph or MGGraph
+        The input graph, for either Single or Multi-GPU operations.
+    
+    edge_biases: FIXE: update this - Create edge_biases of type
+        'cugraph_edge_property_view_t' - edge biases not yet supported.
+
+    start_vertex_list: device array type
+        Device array containing the list of starting vertices for sampling.
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    h_fan_out: tuple of numpy array type
+        Device array containing the branching out (fan-out) degrees per
+        starting vertex for each hop level
+        
+        The sampling method can use different fan_out values for each edge type
+        which is not the case for homogeneous neighborhood sampling.
+
+    with_replacement: bool
+        If true, sampling procedure is done with replacement (the same vertex
+        can be selected multiple times in the same step).
+
+    do_expensive_check: bool
+        If True, performs more extensive tests on the inputs to ensure
+        validitity, at the expense of increased run time.
+
+    with_edge_properties: bool
+        If True, returns the edge properties of each edges along with the
+        edges themselves.  Will result in an error if the provided graph
+        does not have edge properties.
+
+    prior_sources_behavior: str (Optional)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+
+    deduplicate_sources: bool (Optional)
+        If True, will deduplicate the source list before sampling.
+        Defaults to False.
+
+    renumber: bool (Optional)
+        If True, will renumber the sources and destinations on a
+        per-batch basis and return the renumber map and batch offsets
+        in additional to the standard returns.
+
+    retain_seeds: bool (Optional)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+        Defaults to False.
+
+    compression: str (Optional)
+        Options: COO (default), CSR, CSC, DCSR, DCSR
+        Sets the compression format for the returned samples.
+
+    compress_per_hop: bool (Optional)
+        If False (default), will create a compressed edgelist for the
+        entire batch.
+        If True, will create a separate compressed edgelist per hop within
+        a batch.
+
+    random_state: int (Optional)
+        Random state to use when generating samples.  Optional argument,
+        defaults to a hash of process id, time, and hostname.
+        (See pylibcugraph.random.CuGraphRandomState)
+
+    return_dict: bool (Optional)
+        Whether to return a dictionary instead of a tuple.
+        Optional argument, defaults to False, returning a tuple.
+        This argument will eventually be deprecated in favor
+        of always returning a dictionary.
+
+    Returns
+    -------
+    A tuple of device arrays, where the first and second items in the tuple
+    are device arrays containing the starting and ending vertices of each
+    walk respectively, the third item in the tuple is a device array
+    containing the start labels, and the fourth item in the tuple is a device
+    array containing the indices for reconstructing paths.
+
+    If renumber was set to True, then the fifth item in the tuple is a device
+    array containing the renumber map, and the sixth item in the tuple is a
+    device array containing the renumber map offsets (which delineate where
+    the renumber map for each batch starts).
+
+    """
+    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
+        resource_handle.c_resource_handle_ptr
+    )
+
+    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+
+    cdef bool_t c_deduplicate_sources = deduplicate_sources
+    cdef bool_t c_return_hops = return_hops
+    cdef bool_t c_renumber = renumber
+    cdef bool_t c_compress_per_hop = compress_per_hop
+
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+
+    assert_CAI_type(start_vertex_list, "start_vertex_list")
+    assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
+
+    assert_AI_type(h_fan_out, "h_fan_out")
+
+    ai_fan_out_ptr = \
+        h_fan_out.__array_interface__["data"][0]
+
+    fan_out_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_ptr,
+            len(h_fan_out),
+            get_c_type_from_numpy_type(h_fan_out.dtype))
+
+
+
+    cdef cugraph_sample_result_t* result_ptr
+
+    cdef uintptr_t cai_start_ptr = \
+        start_vertex_list.__cuda_array_interface__["data"][0]
+
+    cdef uintptr_t cai_start_vertex_offsets_ptr
+    if start_vertex_offsets is not None:
+        cai_start_vertex_offsets_ptr = \
+            start_vertex_offsets.__cuda_array_interface__['data'][0]
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
+        cugraph_type_erased_device_array_view_create(
+            <void*>cai_start_ptr,
+            len(start_vertex_list),
+            get_c_type_from_numpy_type(start_vertex_list.dtype))
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if start_vertex_offsets is not None:
+        start_vertex_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_start_vertex_offsets_ptr,
+                len(start_vertex_offsets),
+                get_c_type_from_numpy_type(start_vertex_offsets.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if retain_seeds:
+        if start_vertex_offsets is None:
+            raise ValueError("Must provide label offsets if retain_seeds is True")
+
+    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
+
+    cdef cugraph_rng_state_t* rng_state_ptr = \
+        cg_rng_state.rng_state_ptr
+
+    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
+    if prior_sources_behavior is None:
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
+    elif prior_sources_behavior == 'carryover':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
+    elif prior_sources_behavior == 'exclude':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
+    else:
+        raise ValueError(
+            f'Invalid option {prior_sources_behavior}'
+            ' for prior sources behavior'
+        )
+
+    cdef cugraph_compression_type_t compression_behavior_e
+    if compression is None or compression == 'COO':
+        compression_behavior_e = cugraph_compression_type_t.COO
+    elif compression == 'CSR':
+        compression_behavior_e = cugraph_compression_type_t.CSR
+    elif compression == 'CSC':
+        compression_behavior_e = cugraph_compression_type_t.CSC
+    elif compression == 'DCSR':
+        compression_behavior_e = cugraph_compression_type_t.DCSR
+    elif compression == 'DCSC':
+        compression_behavior_e = cugraph_compression_type_t.DCSC
+    else:
+        raise ValueError(
+            f'Invalid option {compression}'
+            ' for compression type'
+        )
+
+    cdef cugraph_sampling_options_t* sampling_options
+    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
+
+    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
+    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
+    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
+    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
+    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
+    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
+    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
+    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
+
+    error_code = cugraph_homogeneous_biased_neighbor_sample(
+        c_resource_handle_ptr,
+        rng_state_ptr,
+        c_graph_ptr,
+        <cugraph_edge_property_view_t*>NULL, # FIXME: Add support for biased neighbor sampling
+        start_vertex_list_ptr,
+        start_vertex_offsets_ptr,
+        fan_out_ptr,
+        sampling_options,
+        do_expensive_check,
+        &result_ptr,
+        &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_homogeneous_biased_neighbor_sample")
+
+    # Free the sampling options
+    cugraph_sampling_options_free(sampling_options)
+
+    # Free the two input arrays that are no longer needed.
+    cugraph_type_erased_device_array_view_free(start_vertex_list_ptr)
+    cugraph_type_erased_host_array_view_free(fan_out_ptr)
+
+    if start_vertex_offsets is not None:
+        cugraph_type_erased_device_array_view_free(start_vertex_offsets_ptr)
+
+    # Have the SamplingResult instance assume ownership of the result data.
+    result = SamplingResult()
+    result.set_ptr(result_ptr)
+
+    # Get cupy "views" of the individual arrays to return. These each increment
+    # the refcount on the SamplingResult instance which will keep the data alive
+    # until all references are removed and the GC runs.
+    # TODO Return everything that isn't null in release 23.12
+    if with_edge_properties:
+        cupy_majors = result.get_majors()
+        cupy_major_offsets = result.get_major_offsets()
+        cupy_minors = result.get_minors()
+        cupy_edge_weights = result.get_edge_weights()
+        cupy_edge_ids = result.get_edge_ids()
+        cupy_edge_types = result.get_edge_types()
+        cupy_batch_ids = result.get_batch_ids()
+        cupy_label_hop_offsets = result.get_label_hop_offsets()
+
+        if renumber:
+            cupy_renumber_map = result.get_renumber_map()
+            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+            # TODO drop the placeholder for hop ids in release 23.12
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': None,
+                    'renumber_map': cupy_renumber_map,
+                    'renumber_map_offsets': cupy_renumber_map_offsets
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+        else:
+            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': cupy_hop_ids,
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
+
+    else:
+        # TODO this is deprecated, remove it in release 23.12
+        warnings.warn(
+            "Calling homogeneous_biased_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
+            "with_edge_properties=True option, returning whatever properties are in the graph.",
+            FutureWarning,
+        )
+
+        cupy_sources = result.get_sources()
+        cupy_destinations = result.get_destinations()
+        cupy_indices = result.get_indices()
+
+        if return_dict:
+            return {
+                'sources': cupy_sources,
+                'destinations': cupy_destinations,
+                'indices': cupy_indices
+            }
+        else:
+            return (cupy_sources, cupy_destinations, cupy_indices)
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
new file mode 100644
index 00000000000..a68e52716e9
--- /dev/null
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
@@ -0,0 +1,411 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Have cython use python 3 syntax
+# cython: language_level = 3
+
+from libc.stdint cimport uintptr_t
+
+from pylibcugraph._cugraph_c.resource_handle cimport (
+    bool_t,
+    cugraph_resource_handle_t,
+)
+from pylibcugraph._cugraph_c.error cimport (
+    cugraph_error_code_t,
+    cugraph_error_t,
+)
+from pylibcugraph._cugraph_c.array cimport (
+    cugraph_type_erased_device_array_view_t,
+    cugraph_type_erased_device_array_view_create,
+    cugraph_type_erased_device_array_view_free,
+    cugraph_type_erased_host_array_view_t,
+    cugraph_type_erased_host_array_view_create,
+    cugraph_type_erased_host_array_view_free,
+)
+from pylibcugraph._cugraph_c.graph cimport (
+    cugraph_graph_t,
+)
+from pylibcugraph._cugraph_c.algorithms cimport (
+    cugraph_sample_result_t,
+    cugraph_prior_sources_behavior_t,
+    cugraph_compression_type_t,
+    cugraph_sampling_options_t,
+    cugraph_sampling_options_create,
+    cugraph_sampling_options_free,
+    cugraph_sampling_set_with_replacement,
+    cugraph_sampling_set_return_hops,
+    cugraph_sampling_set_prior_sources_behavior,
+    cugraph_sampling_set_dedupe_sources,
+    cugraph_sampling_set_renumber_results,
+    cugraph_sampling_set_compress_per_hop,
+    cugraph_sampling_set_compression_type,
+    cugraph_sampling_set_retain_seeds,
+)
+from pylibcugraph._cugraph_c.sampling_algorithms cimport (
+    cugraph_homogeneous_uniform_neighbor_sample,
+)
+from pylibcugraph.resource_handle cimport (
+    ResourceHandle,
+)
+from pylibcugraph.graphs cimport (
+    _GPUGraph,
+)
+from pylibcugraph.utils cimport (
+    assert_success,
+    assert_CAI_type,
+    assert_AI_type,
+    get_c_type_from_numpy_type,
+)
+from pylibcugraph.internal_types.sampling_result cimport (
+    SamplingResult,
+)
+from pylibcugraph._cugraph_c.random cimport (
+    cugraph_rng_state_t
+)
+from pylibcugraph.random cimport (
+    CuGraphRandomState
+)
+import warnings
+
+# TODO accept cupy/numpy random state in addition to raw seed.
+def homogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
+                                          _GPUGraph input_graph,
+                                          start_vertex_list,
+                                          start_vertex_offsets,
+                                          h_fan_out,
+                                          num_edge_types,
+                                          bool_t with_replacement,
+                                          bool_t do_expensive_check,
+                                          with_edge_properties=False,
+                                          prior_sources_behavior=None,
+                                          deduplicate_sources=False,
+                                          return_hops=False,
+                                          renumber=False,
+                                          retain_seeds=False,
+                                          compression='COO',
+                                          compress_per_hop=False,
+                                          random_state=None,
+                                          return_dict=False,):
+    """
+    # FIXME: Deprecate uniform_neighbor_sample
+    Performs biased neighborhood sampling, which samples nodes from
+    a graph based on the current node's neighbors, with a corresponding fan_out
+    value at each hop. The edges are sampled uniformly.
+
+    Parameters
+    ----------
+    resource_handle: ResourceHandle
+        Handle to the underlying device and host resources needed for
+        referencing data and running algorithms.
+
+    input_graph : SGGraph or MGGraph
+        The input graph, for either Single or Multi-GPU operations.
+
+    start_vertex_list: device array type
+        Device array containing the list of starting vertices for sampling.
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    h_fan_out: tuple of numpy array type
+        Device array containing the branching out (fan-out) degrees per
+        starting vertex for each hop level
+        
+        The sampling method can use different fan_out values for each edge type
+        which is not the case for homogeneous neighborhood sampling.
+
+    with_replacement: bool
+        If true, sampling procedure is done with replacement (the same vertex
+        can be selected multiple times in the same step).
+
+    do_expensive_check: bool
+        If True, performs more extensive tests on the inputs to ensure
+        validitity, at the expense of increased run time.
+
+    with_edge_properties: bool
+        If True, returns the edge properties of each edges along with the
+        edges themselves.  Will result in an error if the provided graph
+        does not have edge properties.
+
+    prior_sources_behavior: str (Optional)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+
+    deduplicate_sources: bool (Optional)
+        If True, will deduplicate the source list before sampling.
+        Defaults to False.
+
+    renumber: bool (Optional)
+        If True, will renumber the sources and destinations on a
+        per-batch basis and return the renumber map and batch offsets
+        in additional to the standard returns.
+
+    retain_seeds: bool (Optional)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+        Defaults to False.
+
+    compression: str (Optional)
+        Options: COO (default), CSR, CSC, DCSR, DCSR
+        Sets the compression format for the returned samples.
+
+    compress_per_hop: bool (Optional)
+        If False (default), will create a compressed edgelist for the
+        entire batch.
+        If True, will create a separate compressed edgelist per hop within
+        a batch.
+
+    random_state: int (Optional)
+        Random state to use when generating samples.  Optional argument,
+        defaults to a hash of process id, time, and hostname.
+        (See pylibcugraph.random.CuGraphRandomState)
+
+    return_dict: bool (Optional)
+        Whether to return a dictionary instead of a tuple.
+        Optional argument, defaults to False, returning a tuple.
+        This argument will eventually be deprecated in favor
+        of always returning a dictionary.
+
+    Returns
+    -------
+    A tuple of device arrays, where the first and second items in the tuple
+    are device arrays containing the starting and ending vertices of each
+    walk respectively, the third item in the tuple is a device array
+    containing the start labels, and the fourth item in the tuple is a device
+    array containing the indices for reconstructing paths.
+
+    If renumber was set to True, then the fifth item in the tuple is a device
+    array containing the renumber map, and the sixth item in the tuple is a
+    device array containing the renumber map offsets (which delineate where
+    the renumber map for each batch starts).
+
+    """
+    cdef cugraph_resource_handle_t* c_resource_handle_ptr = (
+        resource_handle.c_resource_handle_ptr
+    )
+
+    cdef cugraph_graph_t* c_graph_ptr = input_graph.c_graph_ptr
+    cdef cugraph_type_erased_host_array_view_t* fan_out_ptr = <cugraph_type_erased_host_array_view_t*>NULL
+
+    cdef bool_t c_deduplicate_sources = deduplicate_sources
+    cdef bool_t c_return_hops = return_hops
+    cdef bool_t c_renumber = renumber
+    cdef bool_t c_compress_per_hop = compress_per_hop
+
+    cdef cugraph_error_code_t error_code
+    cdef cugraph_error_t* error_ptr
+    cdef uintptr_t ai_fan_out_ptr
+
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+
+    assert_CAI_type(start_vertex_list, "start_vertex_list")
+    assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
+
+    assert_AI_type(h_fan_out, "h_fan_out")
+
+    ai_fan_out_ptr = \
+        h_fan_out.__array_interface__["data"][0]
+
+    fan_out_ptr = \
+        cugraph_type_erased_host_array_view_create(
+            <void*>ai_fan_out_ptr,
+            len(h_fan_out),
+            get_c_type_from_numpy_type(h_fan_out.dtype))
+
+
+
+    cdef cugraph_sample_result_t* result_ptr
+
+    cdef uintptr_t cai_start_ptr = \
+        start_vertex_list.__cuda_array_interface__["data"][0]
+
+    cdef uintptr_t cai_start_vertex_offsets_ptr
+    if start_vertex_offsets is not None:
+        cai_start_vertex_offsets_ptr = \
+            start_vertex_offsets.__cuda_array_interface__['data'][0]
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_list_ptr = \
+        cugraph_type_erased_device_array_view_create(
+            <void*>cai_start_ptr,
+            len(start_vertex_list),
+            get_c_type_from_numpy_type(start_vertex_list.dtype))
+
+
+    cdef cugraph_type_erased_device_array_view_t* start_vertex_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if start_vertex_offsets is not None:
+        start_vertex_offsets_ptr = \
+            cugraph_type_erased_device_array_view_create(
+                <void*>cai_start_vertex_offsets_ptr,
+                len(start_vertex_offsets),
+                get_c_type_from_numpy_type(start_vertex_offsets.dtype)
+            )
+
+    cdef cugraph_type_erased_device_array_view_t* label_offsets_ptr = <cugraph_type_erased_device_array_view_t*>NULL
+    if retain_seeds:
+        if start_vertex_offsets is None:
+            raise ValueError("Must provide label offsets if retain_seeds is True")
+
+    cg_rng_state = CuGraphRandomState(resource_handle, random_state)
+
+    cdef cugraph_rng_state_t* rng_state_ptr = \
+        cg_rng_state.rng_state_ptr
+
+    cdef cugraph_prior_sources_behavior_t prior_sources_behavior_e
+    if prior_sources_behavior is None:
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.DEFAULT
+    elif prior_sources_behavior == 'carryover':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.CARRY_OVER
+    elif prior_sources_behavior == 'exclude':
+        prior_sources_behavior_e = cugraph_prior_sources_behavior_t.EXCLUDE
+    else:
+        raise ValueError(
+            f'Invalid option {prior_sources_behavior}'
+            ' for prior sources behavior'
+        )
+
+    cdef cugraph_compression_type_t compression_behavior_e
+    if compression is None or compression == 'COO':
+        compression_behavior_e = cugraph_compression_type_t.COO
+    elif compression == 'CSR':
+        compression_behavior_e = cugraph_compression_type_t.CSR
+    elif compression == 'CSC':
+        compression_behavior_e = cugraph_compression_type_t.CSC
+    elif compression == 'DCSR':
+        compression_behavior_e = cugraph_compression_type_t.DCSR
+    elif compression == 'DCSC':
+        compression_behavior_e = cugraph_compression_type_t.DCSC
+    else:
+        raise ValueError(
+            f'Invalid option {compression}'
+            ' for compression type'
+        )
+
+    cdef cugraph_sampling_options_t* sampling_options
+    error_code = cugraph_sampling_options_create(&sampling_options, &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_sampling_options_create")
+
+    cugraph_sampling_set_with_replacement(sampling_options, with_replacement)
+    cugraph_sampling_set_return_hops(sampling_options, c_return_hops)
+    cugraph_sampling_set_dedupe_sources(sampling_options, c_deduplicate_sources)
+    cugraph_sampling_set_prior_sources_behavior(sampling_options, prior_sources_behavior_e)
+    cugraph_sampling_set_renumber_results(sampling_options, c_renumber)
+    cugraph_sampling_set_compression_type(sampling_options, compression_behavior_e)
+    cugraph_sampling_set_compress_per_hop(sampling_options, c_compress_per_hop)
+    cugraph_sampling_set_retain_seeds(sampling_options, retain_seeds)
+
+    error_code = cugraph_homogeneous_uniform_neighbor_sample(
+        c_resource_handle_ptr,
+        rng_state_ptr,
+        c_graph_ptr,
+        start_vertex_list_ptr,
+        start_vertex_offsets_ptr,
+        fan_out_ptr,
+        sampling_options,
+        do_expensive_check,
+        &result_ptr,
+        &error_ptr)
+    assert_success(error_code, error_ptr, "cugraph_homogeneous_uniform_neighbor_sample")
+
+    # Free the sampling options
+    cugraph_sampling_options_free(sampling_options)
+
+    # Free the two input arrays that are no longer needed.
+    cugraph_type_erased_device_array_view_free(start_vertex_list_ptr)
+    cugraph_type_erased_host_array_view_free(fan_out_ptr)
+
+    if start_vertex_offsets is not None:
+        cugraph_type_erased_device_array_view_free(start_vertex_offsets_ptr)
+
+    # Have the SamplingResult instance assume ownership of the result data.
+    result = SamplingResult()
+    result.set_ptr(result_ptr)
+
+    # Get cupy "views" of the individual arrays to return. These each increment
+    # the refcount on the SamplingResult instance which will keep the data alive
+    # until all references are removed and the GC runs.
+    # TODO Return everything that isn't null in release 23.12
+    if with_edge_properties:
+        cupy_majors = result.get_majors()
+        cupy_major_offsets = result.get_major_offsets()
+        cupy_minors = result.get_minors()
+        cupy_edge_weights = result.get_edge_weights()
+        cupy_edge_ids = result.get_edge_ids()
+        cupy_edge_types = result.get_edge_types()
+        cupy_batch_ids = result.get_batch_ids()
+        cupy_label_hop_offsets = result.get_label_hop_offsets()
+
+        if renumber:
+            cupy_renumber_map = result.get_renumber_map()
+            cupy_renumber_map_offsets = result.get_renumber_map_offsets()
+            # TODO drop the placeholder for hop ids in release 23.12
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': None,
+                    'renumber_map': cupy_renumber_map,
+                    'renumber_map_offsets': cupy_renumber_map_offsets
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, None, cupy_renumber_map, cupy_renumber_map_offsets)
+        else:
+            cupy_hop_ids = result.get_hop_ids() # FIXME remove this
+            if return_dict:
+                return {
+                    'major_offsets': cupy_major_offsets,
+                    'majors': cupy_majors,
+                    'minors': cupy_minors,
+                    'weight': cupy_edge_weights,
+                    'edge_id': cupy_edge_ids,
+                    'edge_type': cupy_edge_types,
+                    'batch_id': cupy_batch_ids,
+                    'label_hop_offsets': cupy_label_hop_offsets,
+                    'hop_id': cupy_hop_ids,
+                }
+            else:
+                cupy_majors = cupy_major_offsets if cupy_majors is None else cupy_majors
+                return (cupy_majors, cupy_minors, cupy_edge_weights, cupy_edge_ids, cupy_edge_types, cupy_batch_ids, cupy_label_hop_offsets, cupy_hop_ids)
+
+    else:
+        # TODO this is deprecated, remove it in release 23.12
+        warnings.warn(
+            "Calling homogeneous_uniform_neighbor_sample with the 'with_edge_properties' argument is deprecated."
+            " Starting in release 23.12, this argument will be removed in favor of behaving like the "
+            "with_edge_properties=True option, returning whatever properties are in the graph.",
+            FutureWarning,
+        )
+
+        cupy_sources = result.get_sources()
+        cupy_destinations = result.get_destinations()
+        cupy_indices = result.get_indices()
+
+        if return_dict:
+            return {
+                'sources': cupy_sources,
+                'destinations': cupy_destinations,
+                'indices': cupy_indices
+            }
+        else:
+            return (cupy_sources, cupy_destinations, cupy_indices)

From f6c4ce31f5fd2ddf4d752e50ad3fc7fdc090713b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 14:14:40 -0700
Subject: [PATCH 060/117] properly handle edge types

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index a5b0ca05a10..d2b74de67b3 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -72,9 +72,6 @@ neighbor_sample_impl(raft::handle_t const& handle,
 {
   static_assert(std::is_floating_point_v<bias_t>);
 
-  std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
-  graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
-
   if constexpr (!multi_gpu) {
     CUGRAPH_EXPECTS(!label_to_output_comm_rank,
                     "cannot specify output GPU mapping in SG implementation");
@@ -103,9 +100,12 @@ neighbor_sample_impl(raft::handle_t const& handle,
     fan_out.size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
     "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                               // to store hops
+  
+  std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
+  graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
   edge_masks_vector.reserve(num_edge_types);
   
-  for (int i = 0; i < num_edge_types; i++) {
+  for (int i = 0; i < num_edge_types - 1; i++) {
 
     cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
     
@@ -170,9 +170,9 @@ neighbor_sample_impl(raft::handle_t const& handle,
   std::vector<size_t> level_sizes{};
 
   // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1
-  auto num_hops = (fan_out.size() % num_edge_types) ? (fan_out.size() / num_edge_types) : (fan_out.size() / num_edge_types + 1);
+  auto num_hops = ((fan_out.size() % num_edge_types) == 0) ? (fan_out.size() / num_edge_types) : ((fan_out.size() / num_edge_types) + 1);
 
-  for (auto hop = 0; hop < num_hops; hop++){
+  for (auto hop = 0; hop < num_hops - 1; hop++){
     for (auto edge_type_id = 0; edge_type_id < num_edge_types; edge_type_id++) {
       auto k_level = fan_out[(hop*num_edge_types) + edge_type_id];
       rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
@@ -182,7 +182,9 @@ neighbor_sample_impl(raft::handle_t const& handle,
       std::optional<rmm::device_uvector<edge_type_t>> edge_types{std::nullopt};
       std::optional<rmm::device_uvector<int32_t>> labels{std::nullopt};
 
-      modified_graph_view.attach_edge_mask(edge_masks_vector[edge_type_id].view());
+      if (num_edge_types > 1) {
+        modified_graph_view.attach_edge_mask(edge_masks_vector[edge_type_id].view());
+      }
 
       if (k_level > 0) {
         std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =

From e71660d610b4aa4393eeebdff2fef19a2ade4bfb Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 14:15:20 -0700
Subject: [PATCH 061/117] add tests for 'homogeneous_uniform_neighbor_sampling'

---
 cpp/tests/CMakeLists.txt                      |   4 +
 .../homogeneous_uniform_neighbor_sampling.cpp | 312 ++++++++++++++++++
 2 files changed, 316 insertions(+)
 create mode 100644 cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index 09b1431e33b..f67fe54615b 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -481,6 +481,10 @@ ConfigureTest(RANDOM_WALKS_TEST sampling/sg_random_walks_test.cpp)
 # - UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
 ConfigureTest(UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/uniform_neighbor_sampling.cpp)
 
+# - UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
+ConfigureTest(
+    HOMOGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/homogeneous_uniform_neighbor_sampling.cpp)
+
 ###################################################################################################
 # - BIASED NBR SAMPLING tests ---------------------------------------------------------------------
 ConfigureTest(BIASED_NEIGHBOR_SAMPLING_TEST sampling/biased_neighbor_sampling.cpp)
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
new file mode 100644
index 00000000000..789ac48111a
--- /dev/null
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -0,0 +1,312 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/property_generator_utilities.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Homogeneous_Uniform_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  bool flag_replacement{true};
+
+  bool edge_masking{false};
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_Homogeneous_Uniform_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_Homogeneous_Uniform_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() {}
+  static void TearDownTestCase() {}
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(
+    std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  {
+    auto [homogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
+
+    raft::handle_t handle{};
+    HighResTimer hr_timer{};
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Construct graph");
+    }
+
+    auto [graph, edge_weights, renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, false>(
+        handle, input_usecase, true, true);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto graph_view = graph.view();
+    auto edge_weight_view =
+      edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
+    if (homogeneous_uniform_neighbor_sampling_usecase.edge_masking) {
+      edge_mask =
+        cugraph::test::generate<decltype(graph_view), bool>::edge_property(handle, graph_view, 2);
+      graph_view.attach_edge_mask((*edge_mask).view());
+    }
+
+    constexpr float select_probability{0.05};
+
+    // FIXME:  Update the tests to initialize RngState and use it instead
+    //         of seed...
+    constexpr uint64_t seed{0};
+
+    raft::random::RngState rng_state(seed);
+
+    auto random_sources = cugraph::select_random_vertices(
+      handle,
+      graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      std::max(static_cast<size_t>(graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(graph_view.number_of_vertices()), size_t{1})),
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
+
+    batch_number = cugraph::test::sequence(
+      handle, random_sources.size(), homogeneous_uniform_neighbor_sampling_usecase.batch_size, int32_t{0});
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle.get_stream());
+
+    std::optional<raft::device_span<int32_t const>>
+      label_to_output_comm_rank_mapping{std::nullopt};
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::homogeneous_uniform_neighbor_sample(
+        handle,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        rng_state,
+        true,
+        homogeneous_uniform_neighbor_sampling_usecase.flag_replacement),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Uniform neighbor sampling");
+    }
+
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::homogeneous_uniform_neighbor_sample(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_uniform_neighbor_sampling_usecase.flag_replacement
+        }
+        );
+    
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (homogeneous_uniform_neighbor_sampling_usecase.check_correctness) {
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * src_out.size(), handle.get_stream());
+      raft::copy(vertices.data(), src_out.data(), src_out.size(), handle.get_stream());
+      raft::copy(
+        vertices.data() + src_out.size(), dst_out.data(), dst_out.size(), handle.get_stream());
+      vertices = cugraph::test::sort<vertex_t>(handle, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(handle, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle.get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle.get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle.get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        handle,
+        graph_view,
+        edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
+        handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
+
+      if (random_sources.size() < 100) {
+        // This validation is too expensive for large number of vertices
+        ASSERT_TRUE(
+          cugraph::test::validate_sampling_depth(handle,
+                                                 std::move(src_out),
+                                                 std::move(dst_out),
+                                                 std::move(wgt_out),
+                                                 std::move(random_sources),
+                                                 homogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
+      }
+    }
+#endif
+  }
+};
+
+using Tests_Homogeneous_Uniform_Neighbor_Sampling_File =
+  Tests_Homogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+using Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat =
+  Tests_Homogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Homogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_Homogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+CUGRAPH_TEST_PROGRAM_MAIN()

From 4e2c8cf38365ac125baf72a84c1665d9aab6ef31 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 15:10:51 -0700
Subject: [PATCH 062/117] add tests for
 homogeneous_biased_neighbor_sampling.cpp

---
 cpp/tests/CMakeLists.txt                      |   7 +-
 .../homogeneous_biased_neighbor_sampling.cpp  | 313 ++++++++++++++++++
 2 files changed, 319 insertions(+), 1 deletion(-)
 create mode 100644 cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index f67fe54615b..7007344261e 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -481,7 +481,7 @@ ConfigureTest(RANDOM_WALKS_TEST sampling/sg_random_walks_test.cpp)
 # - UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
 ConfigureTest(UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/uniform_neighbor_sampling.cpp)
 
-# - UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
+# - HOMOGENEOUS UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
 ConfigureTest(
     HOMOGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/homogeneous_uniform_neighbor_sampling.cpp)
 
@@ -489,6 +489,11 @@ ConfigureTest(
 # - BIASED NBR SAMPLING tests ---------------------------------------------------------------------
 ConfigureTest(BIASED_NEIGHBOR_SAMPLING_TEST sampling/biased_neighbor_sampling.cpp)
 
+###################################################################################################
+# - HOMOGENEOUS BIASED NBR SAMPLING tests ---------------------------------------------------------------------
+ConfigureTest(
+    HOMOGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/homogeneous_biased_neighbor_sampling.cpp)
+
 ###################################################################################################
 # - SAMPLING_POST_PROCESSING tests ----------------------------------------------------------------
 ConfigureTest(SAMPLING_POST_PROCESSING_TEST sampling/sampling_post_processing_test.cpp)
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
new file mode 100644
index 00000000000..61d418e2f4d
--- /dev/null
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -0,0 +1,313 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/property_generator_utilities.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Homogeneous_Biased_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  bool flag_replacement{true};
+
+  bool edge_masking{false};
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_Homogeneous_Biased_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_Homogeneous_Biased_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() {}
+  static void TearDownTestCase() {}
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(
+    std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  {
+    auto [homogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
+
+    raft::handle_t handle{};
+    HighResTimer hr_timer{};
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Construct graph");
+    }
+
+    auto [graph, edge_weights, renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, false>(
+        handle, input_usecase, true, true);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto graph_view = graph.view();
+    auto edge_weight_view =
+      edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
+    if (homogeneous_biased_neighbor_sampling_usecase.edge_masking) {
+      edge_mask =
+        cugraph::test::generate<decltype(graph_view), bool>::edge_property(handle, graph_view, 2);
+      graph_view.attach_edge_mask((*edge_mask).view());
+    }
+
+    constexpr float select_probability{0.05};
+
+    // FIXME:  Update the tests to initialize RngState and use it instead
+    //         of seed...
+    constexpr uint64_t seed{0};
+
+    raft::random::RngState rng_state(seed);
+
+    auto random_sources = cugraph::select_random_vertices(
+      handle,
+      graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      std::max(static_cast<size_t>(graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(graph_view.number_of_vertices()), size_t{1})),
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
+
+    batch_number = cugraph::test::sequence(
+      handle, random_sources.size(), homogeneous_biased_neighbor_sampling_usecase.batch_size, int32_t{0});
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle.get_stream());
+
+    std::optional<raft::device_span<int32_t const>>
+      label_to_output_comm_rank_mapping{std::nullopt};
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::homogeneous_biased_neighbor_sample(
+        handle,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        rng_state,
+        true,
+        homogeneous_biased_neighbor_sampling_usecase.flag_replacement),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Biased neighbor sampling");
+    }
+
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::homogeneous_biased_neighbor_sample(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        *edge_weight_view,
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_biased_neighbor_sampling_usecase.flag_replacement
+        }
+        );
+    
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (homogeneous_biased_neighbor_sampling_usecase.check_correctness) {
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * src_out.size(), handle.get_stream());
+      raft::copy(vertices.data(), src_out.data(), src_out.size(), handle.get_stream());
+      raft::copy(
+        vertices.data() + src_out.size(), dst_out.data(), dst_out.size(), handle.get_stream());
+      vertices = cugraph::test::sort<vertex_t>(handle, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(handle, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle.get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle.get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle.get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        handle,
+        graph_view,
+        edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
+        handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
+
+      if (random_sources.size() < 100) {
+        // This validation is too expensive for large number of vertices
+        ASSERT_TRUE(
+          cugraph::test::validate_sampling_depth(handle,
+                                                 std::move(src_out),
+                                                 std::move(dst_out),
+                                                 std::move(wgt_out),
+                                                 std::move(random_sources),
+                                                 homogeneous_biased_neighbor_sampling_usecase.fanout.size()));
+      }
+    }
+#endif
+  }
+};
+
+using Tests_Homogeneous_Biased_Neighbor_Sampling_File =
+  Tests_Homogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+using Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat =
+  Tests_Homogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Homogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_Homogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+CUGRAPH_TEST_PROGRAM_MAIN()

From 2458149859f6e666c6742d5ba91e51960a8c54da Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 27 Sep 2024 15:13:01 -0700
Subject: [PATCH 063/117] update type combination

---
 .../sampling/neighbor_sampling_mg_v32_e32.cu  | 88 +++++++++++++++++++
 .../sampling/neighbor_sampling_mg_v32_e64.cu  | 88 +++++++++++++++++++
 .../sampling/neighbor_sampling_mg_v64_e64.cu  | 88 +++++++++++++++++++
 .../sampling/neighbor_sampling_sg_v32_e32.cu  | 88 +++++++++++++++++++
 .../sampling/neighbor_sampling_sg_v32_e64.cu  | 88 +++++++++++++++++++
 .../sampling/neighbor_sampling_sg_v64_e64.cu  | 88 +++++++++++++++++++
 6 files changed, 528 insertions(+)

diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
index ea2444c1e3c..b96a73b722d 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
@@ -149,6 +149,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<double>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -172,6 +194,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int32_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                   rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<float>>,
@@ -193,6 +238,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<double>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -215,6 +281,28 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int32_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 
 
 
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
index 846fb33eb21..2a387405900 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
@@ -150,6 +150,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -173,6 +195,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<float>>,
@@ -194,6 +239,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -216,4 +282,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
index 72eec5f6782..505deec51f5 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v64_e64.cu
@@ -149,6 +149,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -172,6 +194,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<float>>,
@@ -193,6 +238,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -215,4 +281,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, true> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
index 46db410f5ae..5cb3fd75ce6 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
@@ -149,6 +149,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<double>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -172,6 +194,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int32_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                   rmm::device_uvector<int32_t>,
                   std::optional<rmm::device_uvector<float>>,
@@ -193,6 +238,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                  rmm::device_uvector<int32_t>,
+                  std::optional<rmm::device_uvector<double>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<int32_t>>,
+                  std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -215,4 +281,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int32_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int32_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int32_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int32_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
index 93c0fe176f4..47cb064cd37 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
@@ -149,6 +149,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -172,6 +194,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<float>>,
@@ -193,6 +238,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int32_t>,
                     rmm::device_uvector<int32_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -215,4 +281,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int32_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
index c3764aa0852..6aa8c71429a 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v64_e64.cu
@@ -149,6 +149,28 @@ heterogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -172,6 +194,29 @@ heterogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+heterogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  int32_t num_edge_types,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<float>>,
@@ -193,6 +238,27 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_uniform_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 template std::tuple<rmm::device_uvector<int64_t>,
                     rmm::device_uvector<int64_t>,
                     std::optional<rmm::device_uvector<double>>,
@@ -215,4 +281,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
+template std::tuple<rmm::device_uvector<int64_t>,
+                    rmm::device_uvector<int64_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+homogeneous_biased_neighbor_sample(
+  raft::handle_t const& handle,
+  raft::random::RngState& rng_state,
+  graph_view_t<int64_t, int64_t, false, false> const& graph_view,
+  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
+  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
+  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
+  edge_property_view_t<int64_t, float const*> edge_bias_view,
+  raft::device_span<int64_t const> starting_vertices,
+  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
+  raft::host_span<int32_t const> fan_out,
+  sampling_flags_t sampling_flags,
+  bool do_expensive_check);
+
 }  // namespace cugraph

From df3e4ffa62dd873f5e9df23c2af7e8f7fc6b90ab Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Sat, 28 Sep 2024 08:28:26 -0700
Subject: [PATCH 064/117] add tests for heterogeneous uniform/biased
 neighborhood sampling

---
 cpp/tests/CMakeLists.txt                      |  13 +-
 ...heterogeneous_biased_neighbor_sampling.cpp | 340 ++++++++++++++++++
 ...eterogeneous_uniform_neighbor_sampling.cpp | 339 +++++++++++++++++
 3 files changed, 690 insertions(+), 2 deletions(-)
 create mode 100644 cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
 create mode 100644 cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index 7007344261e..bfcab9a026b 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -481,19 +481,28 @@ ConfigureTest(RANDOM_WALKS_TEST sampling/sg_random_walks_test.cpp)
 # - UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
 ConfigureTest(UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/uniform_neighbor_sampling.cpp)
 
-# - HOMOGENEOUS UNIFORM NBR SAMPLING tests --------------------------------------------------------------------
+# - HOMOGENEOUS UNIFORM NBR SAMPLING tests --------------------------------------------------------
 ConfigureTest(
     HOMOGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/homogeneous_uniform_neighbor_sampling.cpp)
 
+# - HETEROGENEOUS UNIFORM NBR SAMPLING tests -----------------------------------------------------
+ConfigureTest(
+    HETEROGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/heterogeneous_uniform_neighbor_sampling.cpp)
+
 ###################################################################################################
 # - BIASED NBR SAMPLING tests ---------------------------------------------------------------------
 ConfigureTest(BIASED_NEIGHBOR_SAMPLING_TEST sampling/biased_neighbor_sampling.cpp)
 
 ###################################################################################################
-# - HOMOGENEOUS BIASED NBR SAMPLING tests ---------------------------------------------------------------------
+# - HOMOGENEOUS BIASED NBR SAMPLING tests ---------------------------------------------------------
 ConfigureTest(
     HOMOGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/homogeneous_biased_neighbor_sampling.cpp)
 
+###################################################################################################
+# - HETEROGENEOUS BIASED NBR SAMPLING tests -------------------------------------------------------
+ConfigureTest(
+    HETEROGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/heterogeneous_biased_neighbor_sampling.cpp)
+
 ###################################################################################################
 # - SAMPLING_POST_PROCESSING tests ----------------------------------------------------------------
 ConfigureTest(SAMPLING_POST_PROCESSING_TEST sampling/sampling_post_processing_test.cpp)
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
new file mode 100644
index 00000000000..ca6f1da2ad3
--- /dev/null
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -0,0 +1,340 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/property_generator_utilities.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  int32_t num_edge_types{1};
+  bool flag_replacement{true};
+
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_Heterogeneous_Biased_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_Heterogeneous_Biased_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() {}
+  static void TearDownTestCase() {}
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(
+    std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  {
+    using edge_type_t = int32_t;
+
+    auto [heterogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
+
+    raft::handle_t handle{};
+    HighResTimer hr_timer{};
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Construct graph");
+    }
+
+    auto [graph, edge_weights, renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, false>(
+        handle, input_usecase, true, true);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto graph_view = graph.view();
+    auto edge_weight_view =
+      edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
+
+    constexpr float select_probability{0.05};
+
+    // FIXME:  Update the tests to initialize RngState and use it instead
+    //         of seed...
+    constexpr uint64_t seed{0};
+
+    raft::random::RngState rng_state(seed);
+
+    auto random_sources = cugraph::select_random_vertices(
+      handle,
+      graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      std::max(static_cast<size_t>(graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(graph_view.number_of_vertices()), size_t{1})),
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
+
+    batch_number = cugraph::test::sequence(
+      handle, random_sources.size(), heterogeneous_biased_neighbor_sampling_usecase.batch_size, int32_t{0});
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle.get_stream());
+
+    std::optional<raft::device_span<int32_t const>>
+      label_to_output_comm_rank_mapping{std::nullopt};
+    
+    // Generate the edge types
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), edge_type_t>> edge_types{
+      std::nullopt};
+
+    if (heterogeneous_biased_neighbor_sampling_usecase.num_edge_types > 1) {
+      edge_types = cugraph::test::generate<decltype(graph_view), edge_type_t>::edge_property(
+          handle,
+          graph_view,
+          heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
+    }
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::heterogeneous_biased_neighbor_sample(
+        handle,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        rng_state,
+        true,
+        heterogeneous_biased_neighbor_sampling_usecase.flag_replacement),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Biased neighbor sampling");
+    }
+
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::heterogeneous_biased_neighbor_sample(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        *edge_weight_view,
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_biased_neighbor_sampling_usecase.flag_replacement
+        }
+        );
+    
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (heterogeneous_biased_neighbor_sampling_usecase.check_correctness) {
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * src_out.size(), handle.get_stream());
+      raft::copy(vertices.data(), src_out.data(), src_out.size(), handle.get_stream());
+      raft::copy(
+        vertices.data() + src_out.size(), dst_out.data(), dst_out.size(), handle.get_stream());
+      vertices = cugraph::test::sort<vertex_t>(handle, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(handle, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle.get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle.get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle.get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        handle,
+        graph_view,
+        edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+      
+      
+      
+
+      ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
+        handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
+
+      if (random_sources.size() < 100) {
+        // This validation is too expensive for large number of vertices
+        ASSERT_TRUE(
+          cugraph::test::validate_sampling_depth(handle,
+                                                 std::move(src_out),
+                                                 std::move(dst_out),
+                                                 std::move(wgt_out),
+                                                 std::move(random_sources),
+                                                 heterogeneous_biased_neighbor_sampling_usecase.fanout.size()));
+      }
+    }
+#endif
+  }
+};
+
+using Tests_Heterogeneous_Biased_Neighbor_Sampling_File =
+  Tests_Heterogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
+//#if 0
+using Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat =
+  Tests_Heterogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+//#endif
+
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+#if 0
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+#endif
+
+//#if 0
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
+//#endif
+
+CUGRAPH_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
new file mode 100644
index 00000000000..0b4cf128975
--- /dev/null
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -0,0 +1,339 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/property_generator_utilities.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  int32_t num_edge_types{1};
+  bool flag_replacement{true};
+
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_Heterogeneous_Uniform_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() {}
+  static void TearDownTestCase() {}
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(
+    std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  {
+    using edge_type_t = int32_t;
+
+    auto [heterogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
+
+    raft::handle_t handle{};
+    HighResTimer hr_timer{};
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Construct graph");
+    }
+
+    auto [graph, edge_weights, renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, false>(
+        handle, input_usecase, true, true);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto graph_view = graph.view();
+    auto edge_weight_view =
+      edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
+
+    constexpr float select_probability{0.05};
+
+    // FIXME:  Update the tests to initialize RngState and use it instead
+    //         of seed...
+    constexpr uint64_t seed{0};
+
+    raft::random::RngState rng_state(seed);
+
+    auto random_sources = cugraph::select_random_vertices(
+      handle,
+      graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      std::max(static_cast<size_t>(graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(graph_view.number_of_vertices()), size_t{1})),
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
+
+    batch_number = cugraph::test::sequence(
+      handle, random_sources.size(), heterogeneous_uniform_neighbor_sampling_usecase.batch_size, int32_t{0});
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle.get_stream());
+
+    std::optional<raft::device_span<int32_t const>>
+      label_to_output_comm_rank_mapping{std::nullopt};
+    
+    // Generate the edge types
+
+    std::optional<cugraph::edge_property_t<decltype(graph_view), edge_type_t>> edge_types{
+      std::nullopt};
+
+    if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
+      edge_types = cugraph::test::generate<decltype(graph_view), edge_type_t>::edge_property(
+          handle,
+          graph_view,
+          heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
+    }
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::heterogeneous_uniform_neighbor_sample(
+        handle,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        rng_state,
+        true,
+        heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.start("Uniform neighbor sampling");
+    }
+
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::heterogeneous_uniform_neighbor_sample(
+        handle,
+        rng_state,
+        graph_view,
+        edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
+                                                                           batch_number->size()})
+                     : std::nullopt,
+        label_to_output_comm_rank_mapping,
+        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement
+        }
+        );
+    
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (heterogeneous_uniform_neighbor_sampling_usecase.check_correctness) {
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * src_out.size(), handle.get_stream());
+      raft::copy(vertices.data(), src_out.data(), src_out.size(), handle.get_stream());
+      raft::copy(
+        vertices.data() + src_out.size(), dst_out.data(), dst_out.size(), handle.get_stream());
+      vertices = cugraph::test::sort<vertex_t>(handle, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(handle, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle.get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle.get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle.get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle.get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        handle,
+        graph_view,
+        edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+      
+      
+      
+
+      ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
+        handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
+
+      if (random_sources.size() < 100) {
+        // This validation is too expensive for large number of vertices
+        ASSERT_TRUE(
+          cugraph::test::validate_sampling_depth(handle,
+                                                 std::move(src_out),
+                                                 std::move(dst_out),
+                                                 std::move(wgt_out),
+                                                 std::move(random_sources),
+                                                 heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
+      }
+    }
+#endif
+  }
+};
+
+using Tests_Heterogeneous_Uniform_Neighbor_Sampling_File =
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
+//#if 0
+using Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat =
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+//#endif
+
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+#if 0
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+#endif
+
+//#if 0
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
+//#endif
+
+CUGRAPH_TEST_PROGRAM_MAIN()

From d4847e4afb7a7265bbd13345be4c5b92c3c2fcff Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Sat, 28 Sep 2024 08:33:12 -0700
Subject: [PATCH 065/117] properly sample with edge types

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 46 +++++++++++----------
 1 file changed, 24 insertions(+), 22 deletions(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index d2b74de67b3..5f72a4d99a8 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -100,30 +100,32 @@ neighbor_sample_impl(raft::handle_t const& handle,
     fan_out.size() <= static_cast<size_t>(std::numeric_limits<int32_t>::max()),
     "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                               // to store hops
-  
+
   std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
   graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
   edge_masks_vector.reserve(num_edge_types);
   
-  for (int i = 0; i < num_edge_types - 1; i++) {
-
-    cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
-    
-    cugraph::fill_edge_property(handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
-
-    cugraph::transform_e(
-      handle,
-      modified_graph_view,
-      cugraph::edge_src_dummy_property_t{}.view(),
-      cugraph::edge_dst_dummy_property_t{}.view(),
-      *edge_type_view,
-      [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, auto edge_type) {
-        return edge_type == valid_edge_type;
-      },
-      edge_mask.mutable_view(),
-      false);
-    
-    edge_masks_vector.push_back(std::move(edge_mask));
+  if (num_edge_types > 1) {
+    for (int i = 0; i < num_edge_types; i++) {
+
+      cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
+      
+      cugraph::fill_edge_property(handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
+
+      cugraph::transform_e(
+        handle,
+        modified_graph_view,
+        cugraph::edge_src_dummy_property_t{}.view(),
+        cugraph::edge_dst_dummy_property_t{}.view(),
+        *edge_type_view,
+        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, /*thrust::nullopt_t*/auto edge_type) {
+          return edge_type == valid_edge_type;
+        },
+        edge_mask.mutable_view(),
+        false);
+      
+      edge_masks_vector.push_back(std::move(edge_mask));
+    }
   }
 
   std::vector<rmm::device_uvector<vertex_t>> level_result_src_vectors{};
@@ -172,7 +174,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
   // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1
   auto num_hops = ((fan_out.size() % num_edge_types) == 0) ? (fan_out.size() / num_edge_types) : ((fan_out.size() / num_edge_types) + 1);
 
-  for (auto hop = 0; hop < num_hops - 1; hop++){
+  for (auto hop = 0; hop < num_hops; hop++){
     for (auto edge_type_id = 0; edge_type_id < num_edge_types; edge_type_id++) {
       auto k_level = fan_out[(hop*num_edge_types) + edge_type_id];
       rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
@@ -211,9 +213,9 @@ neighbor_sample_impl(raft::handle_t const& handle,
       }
 
       level_sizes.push_back(srcs.size());
-
       level_result_src_vectors.push_back(std::move(srcs));
       level_result_dst_vectors.push_back(std::move(dsts));
+
       if (weights) { (*level_result_weight_vectors).push_back(std::move(*weights)); }
       if (edge_ids) { (*level_result_edge_id_vectors).push_back(std::move(*edge_ids)); }
       if (edge_types) { (*level_result_edge_type_vectors).push_back(std::move(*edge_types)); }

From dc2c9ba659c7c8f9ff3b838f3da4104e0402bfd4 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Sat, 28 Sep 2024 08:33:38 -0700
Subject: [PATCH 066/117] remove outdated tests

---
 .../heterogeneous_biased_neighbor_sampling.cpp     | 13 +------------
 .../heterogeneous_uniform_neighbor_sampling.cpp    | 14 ++------------
 2 files changed, 3 insertions(+), 24 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index ca6f1da2ad3..e3fd0b841fb 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -238,10 +238,9 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
 
 using Tests_Heterogeneous_Biased_Neighbor_Sampling_File =
   Tests_Heterogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
-//#if 0
+
 using Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat =
   Tests_Heterogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
-//#endif
 
 TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
@@ -249,16 +248,6 @@ TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-#if 0
-INSTANTIATE_TEST_SUITE_P(
-  file_test,
-  Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
-  ::testing::Combine(
-    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
-    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
-#endif
-
-//#if 0
 TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int64Float)
 {
   run_current_test<int32_t, int64_t, float>(
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 0b4cf128975..53c346226f2 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -237,10 +237,10 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
 using Tests_Heterogeneous_Uniform_Neighbor_Sampling_File =
   Tests_Heterogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
-//#if 0
+
 using Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat =
   Tests_Heterogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
-//#endif
+
 
 TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
@@ -248,16 +248,6 @@ TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-#if 0
-INSTANTIATE_TEST_SUITE_P(
-  file_test,
-  Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
-  ::testing::Combine(
-    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
-    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
-#endif
-
-//#if 0
 TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int64Float)
 {
   run_current_test<int32_t, int64_t, float>(

From c01f4e4a8bcf5faa3f6eead4bc2798245cf84e44 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Sun, 29 Sep 2024 22:28:57 -0700
Subject: [PATCH 067/117] add SG python implementation of neighborhood sampling
 both homogeneous and heterogeneous, with/without biases

---
 .../heterogeneous_biased_neighbor_sample.py   | 339 ++++++++++++++++++
 .../heterogeneous_uniform_neighbor_sample.py  | 339 ++++++++++++++++++
 ...ogeneous_biased_uniform_neighbor_sample.py | 333 +++++++++++++++++
 .../homogeneous_uniform_neighbor_sample.py    | 332 +++++++++++++++++
 python/pylibcugraph/pylibcugraph/__init__.py  |   4 +
 5 files changed, 1347 insertions(+)
 create mode 100644 python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
 create mode 100644 python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
 create mode 100644 python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
 create mode 100644 python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py

diff --git a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
new file mode 100644
index 00000000000..b953a5d14f4
--- /dev/null
+++ b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
@@ -0,0 +1,339 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pylibcugraph import ResourceHandle
+from pylibcugraph import heterogeneous_biased_neighbor_sample as \
+    pylibcugraph_heterogeneous_biased_neighbor_sample
+
+from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
+
+import numpy
+
+import cudf
+import cupy as cp
+import warnings
+
+from typing import Union, Tuple, Sequence, List
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cugraph import Graph
+
+
+start_col_name = "_START_"
+batch_col_name = "_BATCH_"
+
+
+# FIXME: Move this function to the utility module so that it can be
+# shared by other algos
+def ensure_valid_dtype(input_graph, start_vertex_list):
+    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
+    if isinstance(start_vertex_list, cudf.Series):
+        start_vertex_list_dtypes = start_vertex_list.dtype
+    else:
+        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
+
+    if start_vertex_list_dtypes != vertex_dtype:
+        warning_msg = (
+            "biased neighbor sample requires 'start_vertex_list' to match the graph's "
+            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
+            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
+        )
+        warnings.warn(warning_msg, UserWarning)
+        start_vertex_list = start_vertex_list.astype(vertex_dtype)
+
+    return start_vertex_list
+
+
+def heterogeneous_biased_neighbor_sample(
+    G: Graph,
+    start_vertex_list: Sequence,
+    start_vertex_offsets: Sequence,
+    fanout_vals: List[int],
+    num_edge_types: int,
+    with_replacement: bool = True,
+    with_edge_properties: bool = False,  # deprecated
+    prior_sources_behavior: str = None,
+    deduplicate_sources: bool = False,
+    return_hops: bool = True,
+    renumber: bool = False,
+    retain_seeds: bool = False,
+    compression: str = "COO",
+    compress_per_hop: bool = False,
+    random_state: int = None,
+    return_offsets: bool = False,
+) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
+    """
+    Does neighborhood sampling, which samples nodes from a graph based on the
+    current node's neighbors, with a corresponding fanout value at each hop.
+
+    Parameters
+    ----------
+    G : cugraph.Graph
+        cuGraph graph, which contains connectivity information as dask cudf
+        edge list dataframe
+
+    start_vertex_list : list or cudf.Series (int32)
+        a list of starting vertices for sampling
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    fanout_vals : list (int32)
+        List of branching out (fan-out) degrees per starting vertex for each
+        hop level.
+    
+    num_edge_types: int32
+        Number of edge types where a value of 1 translates to homogeneous neighbor
+        sample whereas a value greater than 1 translates to heterogeneous neighbor
+        sample.
+
+    with_replacement: bool, optional (default=True)
+        Flag to specify if the random sampling is done with replacement
+
+    with_edge_properties: bool, optional (default=False)
+        Deprecated.
+        Flag to specify whether to return edge properties (weight, edge id,
+        edge type, batch id, hop id) with the sampled edges.
+    
+    prior_sources_behavior: str, optional (default=None)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+    
+    deduplicate_sources: bool, optional (default=False)
+        Whether to first deduplicate the list of possible sources
+        from the previous destinations before performing next
+        hop.
+    
+    return_hops: bool, optional (default=True)
+        Whether to return the sampling results with hop ids
+        corresponding to the hop where the edge appeared.
+        Defaults to True.
+    
+    renumber: bool, optional (default=False)
+        Whether to renumber on a per-batch basis.  If True,
+        will return the renumber map and renumber map offsets
+        as an additional dataframe.
+    
+    retain_seeds: bool, optional (default=False)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+    
+    compression: str, optional (default=COO)
+        Sets the compression type for the output minibatches.
+        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
+    
+    compress_per_hop: bool, optional (default=False)
+        Whether to compress globally (default), or to produce a separate
+        compressed edgelist per hop.
+
+    random_state: int, optional
+        Random seed to use when making sampling calls.
+    
+    return_offsets: bool, optional (default=False)
+        Whether to return the sampling results with batch ids
+        included as one dataframe, or to instead return two
+        dataframes, one with sampling results and one with
+        batch ids and their start offsets.
+
+    Returns
+    -------
+    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
+        GPU data frame containing multiple cudf.Series
+
+        If with_edge_properties=False:
+            df['sources']: cudf.Series
+                Contains the source vertices from the sampling result
+            df['destinations']: cudf.Series
+                Contains the destination vertices from the sampling result
+            df['indices']: cudf.Series
+                Contains the indices (edge weights) from the sampling result
+                for path reconstruction
+
+        If with_edge_properties=True:
+            If return_offsets=False:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+
+            If return_offsets=True:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+
+                offsets_df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                offsets_df['offsets']: cudf.Series
+                    Contains the offsets of each batch in the sampling result
+
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+    """
+
+    use_legacy_names = False # Deprecated parameter
+    include_hop_column=not return_offsets # Deprecated parameter
+
+    major_col_name = "majors"
+    minor_col_name = "minors"
+
+    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
+        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
+
+    if (
+        (compression != "COO")
+        and (not compress_per_hop)
+        and prior_sources_behavior != "exclude"
+    ):
+        raise ValueError(
+            "hop-agnostic compression is only supported with"
+            " the exclude prior sources behavior due to limitations "
+            "of the libcugraph C++ API"
+        )
+
+    if compress_per_hop and prior_sources_behavior != "carryover":
+        raise ValueError(
+            "Compressing the edgelist per hop is only supported "
+            "with the carryover prior sources behavior due to limitations"
+            " of the libcugraph C++ API"
+        )
+
+
+    if with_edge_properties:
+        warning_msg = (
+            "The with_edge_properties flag is deprecated"
+            " and will be removed in the next release in favor"
+            " of returning all properties in the graph"
+        )
+        warnings.warn(warning_msg, FutureWarning)
+
+    if isinstance(start_vertex_list, int):
+        start_vertex_list = [start_vertex_list]
+
+    if isinstance(start_vertex_list, list):
+        start_vertex_list = cudf.Series(
+            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
+        )
+
+
+    """ 
+    # No batch_ids, the rank owning the vertices will wom the final
+    # result.
+    if with_edge_properties and not with_batch_ids:
+        if isinstance(start_vertex_list, cudf.Series):
+            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
+
+        start_vertex_list[batch_col_name] = cudf.Series(
+            cp.zeros(len(start_vertex_list), dtype="int32")
+        )
+    """
+
+    # fanout_vals must be passed to pylibcugraph as a host array
+    if isinstance(fanout_vals, numpy.ndarray):
+        fanout_vals = fanout_vals.astype("int32")
+    elif isinstance(fanout_vals, list):
+        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
+    elif isinstance(fanout_vals, cp.ndarray):
+        fanout_vals = fanout_vals.get().astype("int32")
+    elif isinstance(fanout_vals, cudf.Series):
+        fanout_vals = fanout_vals.values_host.astype("int32")
+    else:
+        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
+
+    if "weights" in G.edgelist.edgelist_df:
+        weight_t = G.edgelist.edgelist_df["weights"].dtype
+    else:
+        weight_t = "float32"
+
+    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
+
+
+    if G.renumbered:
+        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
+
+
+    sampling_result_array_dict = pylibcugraph_heterogeneous_biased_neighbor_sample(
+        resource_handle=ResourceHandle(),
+        input_graph=G._plc_graph,
+        start_vertex_list=start_vertex_list,
+        start_vertex_offsets=start_vertex_offsets,
+        h_fan_out=fanout_vals,
+        num_edge_types=num_edge_types,
+        with_replacement=with_replacement,
+        do_expensive_check=False,
+        with_edge_properties=with_edge_properties,
+        prior_sources_behavior=prior_sources_behavior,
+        deduplicate_sources=deduplicate_sources,
+        return_hops=return_hops,
+        renumber=renumber,
+        retain_seeds=retain_seeds,
+        compression=compression,
+        compress_per_hop=compress_per_hop,
+        random_state=random_state,
+        return_dict=True,
+    )
+
+    dfs = sampling_results_from_cupy_array_dict(
+        sampling_result_array_dict,
+        weight_t,
+        len(fanout_vals),
+        with_edge_properties=with_edge_properties,
+        return_offsets=return_offsets,
+        renumber=renumber,
+        use_legacy_names=use_legacy_names,
+        include_hop_column=include_hop_column, # Deprecated flag
+    )
+
+    if G.renumbered and not renumber:
+        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
+        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
+
+    if len(dfs) > 1:
+        return dfs
+
+    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
new file mode 100644
index 00000000000..7eadebbc912
--- /dev/null
+++ b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
@@ -0,0 +1,339 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pylibcugraph import ResourceHandle
+from pylibcugraph import heterogeneous_uniform_neighbor_sample as \
+    pylibcugraph_heterogeneous_uniform_neighbor_sample
+
+from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
+
+import numpy
+
+import cudf
+import cupy as cp
+import warnings
+
+from typing import Union, Tuple, Sequence, List
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cugraph import Graph
+
+
+start_col_name = "_START_"
+batch_col_name = "_BATCH_"
+
+
+# FIXME: Move this function to the utility module so that it can be
+# shared by other algos
+def ensure_valid_dtype(input_graph, start_vertex_list):
+    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
+    if isinstance(start_vertex_list, cudf.Series):
+        start_vertex_list_dtypes = start_vertex_list.dtype
+    else:
+        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
+
+    if start_vertex_list_dtypes != vertex_dtype:
+        warning_msg = (
+            "Uniform neighbor sample requires 'start_vertex_list' to match the graph's "
+            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
+            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
+        )
+        warnings.warn(warning_msg, UserWarning)
+        start_vertex_list = start_vertex_list.astype(vertex_dtype)
+
+    return start_vertex_list
+
+
+def heterogeneous_uniform_neighbor_sample(
+    G: Graph,
+    start_vertex_list: Sequence,
+    start_vertex_offsets: Sequence,
+    fanout_vals: List[int],
+    num_edge_types: int,
+    with_replacement: bool = True,
+    with_edge_properties: bool = False,  # deprecated
+    prior_sources_behavior: str = None,
+    deduplicate_sources: bool = False,
+    return_hops: bool = True,
+    renumber: bool = False,
+    retain_seeds: bool = False,
+    compression: str = "COO",
+    compress_per_hop: bool = False,
+    random_state: int = None,
+    return_offsets: bool = False,
+) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
+    """
+    Does neighborhood sampling, which samples nodes from a graph based on the
+    current node's neighbors, with a corresponding fanout value at each hop.
+
+    Parameters
+    ----------
+    G : cugraph.Graph
+        cuGraph graph, which contains connectivity information as dask cudf
+        edge list dataframe
+
+    start_vertex_list : list or cudf.Series (int32)
+        a list of starting vertices for sampling
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    fanout_vals : list (int32)
+        List of branching out (fan-out) degrees per starting vertex for each
+        hop level.
+    
+    num_edge_types: int32
+        Number of edge types where a value of 1 translates to homogeneous neighbor
+        sample whereas a value greater than 1 translates to heterogeneous neighbor
+        sample.
+
+    with_replacement: bool, optional (default=True)
+        Flag to specify if the random sampling is done with replacement
+
+    with_edge_properties: bool, optional (default=False)
+        Deprecated.
+        Flag to specify whether to return edge properties (weight, edge id,
+        edge type, batch id, hop id) with the sampled edges.
+    
+    prior_sources_behavior: str, optional (default=None)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+    
+    deduplicate_sources: bool, optional (default=False)
+        Whether to first deduplicate the list of possible sources
+        from the previous destinations before performing next
+        hop.
+    
+    return_hops: bool, optional (default=True)
+        Whether to return the sampling results with hop ids
+        corresponding to the hop where the edge appeared.
+        Defaults to True.
+    
+    renumber: bool, optional (default=False)
+        Whether to renumber on a per-batch basis.  If True,
+        will return the renumber map and renumber map offsets
+        as an additional dataframe.
+    
+    retain_seeds: bool, optional (default=False)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+    
+    compression: str, optional (default=COO)
+        Sets the compression type for the output minibatches.
+        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
+    
+    compress_per_hop: bool, optional (default=False)
+        Whether to compress globally (default), or to produce a separate
+        compressed edgelist per hop.
+
+    random_state: int, optional
+        Random seed to use when making sampling calls.
+    
+    return_offsets: bool, optional (default=False)
+        Whether to return the sampling results with batch ids
+        included as one dataframe, or to instead return two
+        dataframes, one with sampling results and one with
+        batch ids and their start offsets.
+
+    Returns
+    -------
+    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
+        GPU data frame containing multiple cudf.Series
+
+        If with_edge_properties=False:
+            df['sources']: cudf.Series
+                Contains the source vertices from the sampling result
+            df['destinations']: cudf.Series
+                Contains the destination vertices from the sampling result
+            df['indices']: cudf.Series
+                Contains the indices (edge weights) from the sampling result
+                for path reconstruction
+
+        If with_edge_properties=True:
+            If return_offsets=False:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+
+            If return_offsets=True:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+
+                offsets_df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                offsets_df['offsets']: cudf.Series
+                    Contains the offsets of each batch in the sampling result
+
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+    """
+
+    use_legacy_names = False # Deprecated parameter
+    include_hop_column=not return_offsets # Deprecated parameter
+
+    major_col_name = "majors"
+    minor_col_name = "minors"
+
+    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
+        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
+
+    if (
+        (compression != "COO")
+        and (not compress_per_hop)
+        and prior_sources_behavior != "exclude"
+    ):
+        raise ValueError(
+            "hop-agnostic compression is only supported with"
+            " the exclude prior sources behavior due to limitations "
+            "of the libcugraph C++ API"
+        )
+
+    if compress_per_hop and prior_sources_behavior != "carryover":
+        raise ValueError(
+            "Compressing the edgelist per hop is only supported "
+            "with the carryover prior sources behavior due to limitations"
+            " of the libcugraph C++ API"
+        )
+
+
+    if with_edge_properties:
+        warning_msg = (
+            "The with_edge_properties flag is deprecated"
+            " and will be removed in the next release in favor"
+            " of returning all properties in the graph"
+        )
+        warnings.warn(warning_msg, FutureWarning)
+
+    if isinstance(start_vertex_list, int):
+        start_vertex_list = [start_vertex_list]
+
+    if isinstance(start_vertex_list, list):
+        start_vertex_list = cudf.Series(
+            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
+        )
+
+
+    """ 
+    # No batch_ids, the rank owning the vertices will wom the final
+    # result.
+    if with_edge_properties and not with_batch_ids:
+        if isinstance(start_vertex_list, cudf.Series):
+            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
+
+        start_vertex_list[batch_col_name] = cudf.Series(
+            cp.zeros(len(start_vertex_list), dtype="int32")
+        )
+    """
+
+    # fanout_vals must be passed to pylibcugraph as a host array
+    if isinstance(fanout_vals, numpy.ndarray):
+        fanout_vals = fanout_vals.astype("int32")
+    elif isinstance(fanout_vals, list):
+        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
+    elif isinstance(fanout_vals, cp.ndarray):
+        fanout_vals = fanout_vals.get().astype("int32")
+    elif isinstance(fanout_vals, cudf.Series):
+        fanout_vals = fanout_vals.values_host.astype("int32")
+    else:
+        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
+
+    if "weights" in G.edgelist.edgelist_df:
+        weight_t = G.edgelist.edgelist_df["weights"].dtype
+    else:
+        weight_t = "float32"
+
+    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
+
+
+    if G.renumbered:
+        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
+
+
+    sampling_result_array_dict = pylibcugraph_heterogeneous_uniform_neighbor_sample(
+        resource_handle=ResourceHandle(),
+        input_graph=G._plc_graph,
+        start_vertex_list=start_vertex_list,
+        start_vertex_offsets=start_vertex_offsets,
+        h_fan_out=fanout_vals,
+        num_edge_types=num_edge_types,
+        with_replacement=with_replacement,
+        do_expensive_check=False,
+        with_edge_properties=with_edge_properties,
+        prior_sources_behavior=prior_sources_behavior,
+        deduplicate_sources=deduplicate_sources,
+        return_hops=return_hops,
+        renumber=renumber,
+        retain_seeds=retain_seeds,
+        compression=compression,
+        compress_per_hop=compress_per_hop,
+        random_state=random_state,
+        return_dict=True,
+    )
+
+    dfs = sampling_results_from_cupy_array_dict(
+        sampling_result_array_dict,
+        weight_t,
+        len(fanout_vals),
+        with_edge_properties=with_edge_properties,
+        return_offsets=return_offsets,
+        renumber=renumber,
+        use_legacy_names=use_legacy_names,
+        include_hop_column=include_hop_column, # Deprecated flag
+    )
+
+    if G.renumbered and not renumber:
+        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
+        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
+
+    if len(dfs) > 1:
+        return dfs
+
+    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
new file mode 100644
index 00000000000..62ea1a93686
--- /dev/null
+++ b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
@@ -0,0 +1,333 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pylibcugraph import ResourceHandle
+from pylibcugraph import homogeneous_biased_neighbor_sample as \
+    pylibcugraph_homogeneous_biased_neighbor_sample
+
+from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
+
+import numpy
+
+import cudf
+import cupy as cp
+import warnings
+
+from typing import Union, Tuple, Sequence, List
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cugraph import Graph
+
+
+start_col_name = "_START_"
+batch_col_name = "_BATCH_"
+
+
+# FIXME: Move this function to the utility module so that it can be
+# shared by other algos
+def ensure_valid_dtype(input_graph, start_vertex_list):
+    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
+    if isinstance(start_vertex_list, cudf.Series):
+        start_vertex_list_dtypes = start_vertex_list.dtype
+    else:
+        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
+
+    if start_vertex_list_dtypes != vertex_dtype:
+        warning_msg = (
+            "biased neighbor sample requires 'start_vertex_list' to match the graph's "
+            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
+            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
+        )
+        warnings.warn(warning_msg, UserWarning)
+        start_vertex_list = start_vertex_list.astype(vertex_dtype)
+
+    return start_vertex_list
+
+
+def homogeneous_biased_neighbor_sample(
+    G: Graph,
+    start_vertex_list: Sequence,
+    start_vertex_offsets: Sequence,
+    fanout_vals: List[int],
+    with_replacement: bool = True,
+    with_edge_properties: bool = False,  # deprecated
+    prior_sources_behavior: str = None,
+    deduplicate_sources: bool = False,
+    return_hops: bool = True,
+    renumber: bool = False,
+    retain_seeds: bool = False,
+    compression: str = "COO",
+    compress_per_hop: bool = False,
+    random_state: int = None,
+    return_offsets: bool = False,
+) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
+    """
+    Does neighborhood sampling, which samples nodes from a graph based on the
+    current node's neighbors, with a corresponding fanout value at each hop.
+    The edges are sample with biases.
+
+    Parameters
+    ----------
+    G : cugraph.Graph
+        cuGraph graph, which contains connectivity information as dask cudf
+        edge list dataframe
+
+    start_vertex_list : list or cudf.Series (int32)
+        a list of starting vertices for sampling
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    fanout_vals : list (int32)
+        List of branching out (fan-out) degrees per starting vertex for each
+        hop level.
+
+    with_replacement: bool, optional (default=True)
+        Flag to specify if the random sampling is done with replacement
+
+    with_edge_properties: bool, optional (default=False)
+        Deprecated.
+        Flag to specify whether to return edge properties (weight, edge id,
+        edge type, batch id, hop id) with the sampled edges.
+    
+    prior_sources_behavior: str, optional (default=None)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+    
+    deduplicate_sources: bool, optional (default=False)
+        Whether to first deduplicate the list of possible sources
+        from the previous destinations before performing next
+        hop.
+    
+    return_hops: bool, optional (default=True)
+        Whether to return the sampling results with hop ids
+        corresponding to the hop where the edge appeared.
+        Defaults to True.
+    
+    renumber: bool, optional (default=False)
+        Whether to renumber on a per-batch basis.  If True,
+        will return the renumber map and renumber map offsets
+        as an additional dataframe.
+    
+    retain_seeds: bool, optional (default=False)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+    
+    compression: str, optional (default=COO)
+        Sets the compression type for the output minibatches.
+        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
+    
+    compress_per_hop: bool, optional (default=False)
+        Whether to compress globally (default), or to produce a separate
+        compressed edgelist per hop.
+
+    random_state: int, optional
+        Random seed to use when making sampling calls.
+    
+    return_offsets: bool, optional (default=False)
+        Whether to return the sampling results with batch ids
+        included as one dataframe, or to instead return two
+        dataframes, one with sampling results and one with
+        batch ids and their start offsets.
+
+    Returns
+    -------
+    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
+        GPU data frame containing multiple cudf.Series
+
+        If with_edge_properties=False:
+            df['sources']: cudf.Series
+                Contains the source vertices from the sampling result
+            df['destinations']: cudf.Series
+                Contains the destination vertices from the sampling result
+            df['indices']: cudf.Series
+                Contains the indices (edge weights) from the sampling result
+                for path reconstruction
+
+        If with_edge_properties=True:
+            If return_offsets=False:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+
+            If return_offsets=True:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+
+                offsets_df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                offsets_df['offsets']: cudf.Series
+                    Contains the offsets of each batch in the sampling result
+
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+    """
+
+    use_legacy_names = False # Deprecated parameter
+    include_hop_column=not return_offsets # Deprecated parameter
+
+    major_col_name = "majors"
+    minor_col_name = "minors"
+
+    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
+        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
+
+    if (
+        (compression != "COO")
+        and (not compress_per_hop)
+        and prior_sources_behavior != "exclude"
+    ):
+        raise ValueError(
+            "hop-agnostic compression is only supported with"
+            " the exclude prior sources behavior due to limitations "
+            "of the libcugraph C++ API"
+        )
+
+    if compress_per_hop and prior_sources_behavior != "carryover":
+        raise ValueError(
+            "Compressing the edgelist per hop is only supported "
+            "with the carryover prior sources behavior due to limitations"
+            " of the libcugraph C++ API"
+        )
+
+
+    if with_edge_properties:
+        warning_msg = (
+            "The with_edge_properties flag is deprecated"
+            " and will be removed in the next release in favor"
+            " of returning all properties in the graph"
+        )
+        warnings.warn(warning_msg, FutureWarning)
+
+    if isinstance(start_vertex_list, int):
+        start_vertex_list = [start_vertex_list]
+
+    if isinstance(start_vertex_list, list):
+        start_vertex_list = cudf.Series(
+            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
+        )
+
+
+    """ 
+    # No batch_ids, the rank owning the vertices will wom the final
+    # result.
+    if with_edge_properties and not with_batch_ids:
+        if isinstance(start_vertex_list, cudf.Series):
+            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
+
+        start_vertex_list[batch_col_name] = cudf.Series(
+            cp.zeros(len(start_vertex_list), dtype="int32")
+        )
+    """
+
+    # fanout_vals must be passed to pylibcugraph as a host array
+    if isinstance(fanout_vals, numpy.ndarray):
+        fanout_vals = fanout_vals.astype("int32")
+    elif isinstance(fanout_vals, list):
+        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
+    elif isinstance(fanout_vals, cp.ndarray):
+        fanout_vals = fanout_vals.get().astype("int32")
+    elif isinstance(fanout_vals, cudf.Series):
+        fanout_vals = fanout_vals.values_host.astype("int32")
+    else:
+        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
+
+    if "weights" in G.edgelist.edgelist_df:
+        weight_t = G.edgelist.edgelist_df["weights"].dtype
+    else:
+        weight_t = "float32"
+
+    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
+
+
+    if G.renumbered:
+        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
+
+
+    sampling_result_array_dict = pylibcugraph_homogeneous_biased_neighbor_sample(
+        resource_handle=ResourceHandle(),
+        input_graph=G._plc_graph,
+        start_vertex_list=start_vertex_list,
+        start_vertex_offsets=start_vertex_offsets,
+        h_fan_out=fanout_vals,
+        with_replacement=with_replacement,
+        do_expensive_check=False,
+        with_edge_properties=with_edge_properties,
+        prior_sources_behavior=prior_sources_behavior,
+        deduplicate_sources=deduplicate_sources,
+        return_hops=return_hops,
+        renumber=renumber,
+        retain_seeds=retain_seeds,
+        compression=compression,
+        compress_per_hop=compress_per_hop,
+        random_state=random_state,
+        return_dict=True,
+    )
+
+    dfs = sampling_results_from_cupy_array_dict(
+        sampling_result_array_dict,
+        weight_t,
+        len(fanout_vals),
+        with_edge_properties=with_edge_properties,
+        return_offsets=return_offsets,
+        renumber=renumber,
+        use_legacy_names=use_legacy_names,
+        include_hop_column=include_hop_column, # Deprecated flag
+    )
+
+    if G.renumbered and not renumber:
+        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
+        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
+
+    if len(dfs) > 1:
+        return dfs
+
+    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
new file mode 100644
index 00000000000..4c015251af7
--- /dev/null
+++ b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
@@ -0,0 +1,332 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pylibcugraph import ResourceHandle
+from pylibcugraph import homogeneous_uniform_neighbor_sample as \
+    pylibcugraph_homogeneous_uniform_neighbor_sample
+
+from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
+
+import numpy
+
+import cudf
+import cupy as cp
+import warnings
+
+from typing import Union, Tuple, Sequence, List
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from cugraph import Graph
+
+
+start_col_name = "_START_"
+batch_col_name = "_BATCH_"
+
+
+# FIXME: Move this function to the utility module so that it can be
+# shared by other algos
+def ensure_valid_dtype(input_graph, start_vertex_list):
+    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
+    if isinstance(start_vertex_list, cudf.Series):
+        start_vertex_list_dtypes = start_vertex_list.dtype
+    else:
+        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
+
+    if start_vertex_list_dtypes != vertex_dtype:
+        warning_msg = (
+            "Uniform neighbor sample requires 'start_vertex_list' to match the graph's "
+            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
+            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
+        )
+        warnings.warn(warning_msg, UserWarning)
+        start_vertex_list = start_vertex_list.astype(vertex_dtype)
+
+    return start_vertex_list
+
+
+def homogeneous_uniform_neighbor_sample(
+    G: Graph,
+    start_vertex_list: Sequence,
+    start_vertex_offsets: Sequence,
+    fanout_vals: List[int],
+    with_replacement: bool = True,
+    with_edge_properties: bool = False,  # deprecated
+    prior_sources_behavior: str = None,
+    deduplicate_sources: bool = False,
+    return_hops: bool = True,
+    renumber: bool = False,
+    retain_seeds: bool = False,
+    compression: str = "COO",
+    compress_per_hop: bool = False,
+    random_state: int = None,
+    return_offsets: bool = False,
+) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
+    """
+    Does neighborhood sampling, which samples nodes from a graph based on the
+    current node's neighbors, with a corresponding fanout value at each hop.
+
+    Parameters
+    ----------
+    G : cugraph.Graph
+        cuGraph graph, which contains connectivity information as dask cudf
+        edge list dataframe
+
+    start_vertex_list : list or cudf.Series (int32)
+        a list of starting vertices for sampling
+    
+    start_vertex_offsets: list[int] (Optional)
+        Offsets of each label within the start vertex list.
+
+    fanout_vals : list (int32)
+        List of branching out (fan-out) degrees per starting vertex for each
+        hop level.
+
+    with_replacement: bool, optional (default=True)
+        Flag to specify if the random sampling is done with replacement
+
+    with_edge_properties: bool, optional (default=False)
+        Deprecated.
+        Flag to specify whether to return edge properties (weight, edge id,
+        edge type, batch id, hop id) with the sampled edges.
+    
+    prior_sources_behavior: str, optional (default=None)
+        Options are "carryover", and "exclude".
+        Default will leave the source list as-is.
+        Carryover will carry over sources from previous hops to the
+        current hop.
+        Exclude will exclude sources from previous hops from reappearing
+        as sources in future hops.
+    
+    deduplicate_sources: bool, optional (default=False)
+        Whether to first deduplicate the list of possible sources
+        from the previous destinations before performing next
+        hop.
+    
+    return_hops: bool, optional (default=True)
+        Whether to return the sampling results with hop ids
+        corresponding to the hop where the edge appeared.
+        Defaults to True.
+    
+    renumber: bool, optional (default=False)
+        Whether to renumber on a per-batch basis.  If True,
+        will return the renumber map and renumber map offsets
+        as an additional dataframe.
+    
+    retain_seeds: bool, optional (default=False)
+        If True, will retain the original seeds (original source vertices)
+        in the output even if they do not have outgoing neighbors.
+    
+    compression: str, optional (default=COO)
+        Sets the compression type for the output minibatches.
+        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
+    
+    compress_per_hop: bool, optional (default=False)
+        Whether to compress globally (default), or to produce a separate
+        compressed edgelist per hop.
+
+    random_state: int, optional
+        Random seed to use when making sampling calls.
+    
+    return_offsets: bool, optional (default=False)
+        Whether to return the sampling results with batch ids
+        included as one dataframe, or to instead return two
+        dataframes, one with sampling results and one with
+        batch ids and their start offsets.
+
+    Returns
+    -------
+    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
+        GPU data frame containing multiple cudf.Series
+
+        If with_edge_properties=False:
+            df['sources']: cudf.Series
+                Contains the source vertices from the sampling result
+            df['destinations']: cudf.Series
+                Contains the destination vertices from the sampling result
+            df['indices']: cudf.Series
+                Contains the indices (edge weights) from the sampling result
+                for path reconstruction
+
+        If with_edge_properties=True:
+            If return_offsets=False:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+
+            If return_offsets=True:
+                df['sources']: cudf.Series
+                    Contains the source vertices from the sampling result
+                df['destinations']: cudf.Series
+                    Contains the destination vertices from the sampling result
+                df['edge_weight']: cudf.Series
+                    Contains the edge weights from the sampling result
+                df['edge_id']: cudf.Series
+                    Contains the edge ids from the sampling result
+                df['edge_type']: cudf.Series
+                    Contains the edge types from the sampling result
+                df['hop_id']: cudf.Series
+                    Contains the hop ids from the sampling result
+
+                offsets_df['batch_id']: cudf.Series
+                    Contains the batch ids from the sampling result
+                offsets_df['offsets']: cudf.Series
+                    Contains the offsets of each batch in the sampling result
+
+                If renumber=True:
+                    (adds the following dataframe)
+                    renumber_df['map']: cudf.Series
+                        Contains the renumber maps for each batch
+                    renumber_df['offsets']: cudf.Series
+                        Contains the batch offsets for the renumber maps
+    """
+
+    use_legacy_names = False # Deprecated parameter
+    include_hop_column=not return_offsets # Deprecated parameter
+
+    major_col_name = "majors"
+    minor_col_name = "minors"
+
+    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
+        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
+
+    if (
+        (compression != "COO")
+        and (not compress_per_hop)
+        and prior_sources_behavior != "exclude"
+    ):
+        raise ValueError(
+            "hop-agnostic compression is only supported with"
+            " the exclude prior sources behavior due to limitations "
+            "of the libcugraph C++ API"
+        )
+
+    if compress_per_hop and prior_sources_behavior != "carryover":
+        raise ValueError(
+            "Compressing the edgelist per hop is only supported "
+            "with the carryover prior sources behavior due to limitations"
+            " of the libcugraph C++ API"
+        )
+
+
+    if with_edge_properties:
+        warning_msg = (
+            "The with_edge_properties flag is deprecated"
+            " and will be removed in the next release in favor"
+            " of returning all properties in the graph"
+        )
+        warnings.warn(warning_msg, FutureWarning)
+
+    if isinstance(start_vertex_list, int):
+        start_vertex_list = [start_vertex_list]
+
+    if isinstance(start_vertex_list, list):
+        start_vertex_list = cudf.Series(
+            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
+        )
+
+
+    """ 
+    # No batch_ids, the rank owning the vertices will wom the final
+    # result.
+    if with_edge_properties and not with_batch_ids:
+        if isinstance(start_vertex_list, cudf.Series):
+            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
+
+        start_vertex_list[batch_col_name] = cudf.Series(
+            cp.zeros(len(start_vertex_list), dtype="int32")
+        )
+    """
+
+    # fanout_vals must be passed to pylibcugraph as a host array
+    if isinstance(fanout_vals, numpy.ndarray):
+        fanout_vals = fanout_vals.astype("int32")
+    elif isinstance(fanout_vals, list):
+        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
+    elif isinstance(fanout_vals, cp.ndarray):
+        fanout_vals = fanout_vals.get().astype("int32")
+    elif isinstance(fanout_vals, cudf.Series):
+        fanout_vals = fanout_vals.values_host.astype("int32")
+    else:
+        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
+
+    if "weights" in G.edgelist.edgelist_df:
+        weight_t = G.edgelist.edgelist_df["weights"].dtype
+    else:
+        weight_t = "float32"
+
+    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
+
+
+    if G.renumbered:
+        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
+
+
+    sampling_result_array_dict = pylibcugraph_homogeneous_uniform_neighbor_sample(
+        resource_handle=ResourceHandle(),
+        input_graph=G._plc_graph,
+        start_vertex_list=start_vertex_list,
+        start_vertex_offsets=start_vertex_offsets,
+        h_fan_out=fanout_vals,
+        with_replacement=with_replacement,
+        do_expensive_check=False,
+        with_edge_properties=with_edge_properties,
+        prior_sources_behavior=prior_sources_behavior,
+        deduplicate_sources=deduplicate_sources,
+        return_hops=return_hops,
+        renumber=renumber,
+        retain_seeds=retain_seeds,
+        compression=compression,
+        compress_per_hop=compress_per_hop,
+        random_state=random_state,
+        return_dict=True,
+    )
+
+    dfs = sampling_results_from_cupy_array_dict(
+        sampling_result_array_dict,
+        weight_t,
+        len(fanout_vals),
+        with_edge_properties=with_edge_properties,
+        return_offsets=return_offsets,
+        renumber=renumber,
+        use_legacy_names=use_legacy_names,
+        include_hop_column=include_hop_column, # Deprecated flag
+    )
+
+    if G.renumbered and not renumber:
+        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
+        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
+
+    if len(dfs) > 1:
+        return dfs
+
+    return dfs[0]
diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py
index 7c241e439c7..09a584c6da4 100644
--- a/python/pylibcugraph/pylibcugraph/__init__.py
+++ b/python/pylibcugraph/pylibcugraph/__init__.py
@@ -38,6 +38,10 @@
 from pylibcugraph.bfs import bfs
 
 from pylibcugraph.uniform_neighbor_sample import uniform_neighbor_sample
+from pylibcugraph.homogeneous_uniform_neighbor_sample import homogeneous_uniform_neighbor_sample
+from pylibcugraph.homogeneous_biased_neighbor_sample import homogeneous_biased_neighbor_sample
+from pylibcugraph.heterogeneous_uniform_neighbor_sample import heterogeneous_uniform_neighbor_sample
+from pylibcugraph.heterogeneous_biased_neighbor_sample import heterogeneous_biased_neighbor_sample
 
 
 #FIXME: break down the API into homogeneous and heterogeneous neighbor sample

From dabd0c8f09428dcd662edca6bb2b6f5d2e266262 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Sun, 29 Sep 2024 22:29:30 -0700
Subject: [PATCH 068/117] remove unused argument

---
 .../pylibcugraph/heterogeneous_biased_neighbor_sample.pyx      | 2 +-
 .../pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx     | 2 +-
 .../pylibcugraph/homogeneous_biased_neighbor_sample.pyx        | 3 +--
 .../pylibcugraph/homogeneous_uniform_neighbor_sample.pyx       | 3 +--
 4 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
index 53de709e28d..401b2ee6258 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
@@ -220,7 +220,7 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
     cdef cugraph_error_t* error_ptr
     cdef uintptr_t ai_fan_out_ptr
 
-    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
index 2bcf0d856fc..477a010037b 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
@@ -214,7 +214,7 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
     cdef cugraph_error_t* error_ptr
     cdef uintptr_t ai_fan_out_ptr
 
-    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
index 926d907995e..d95d8a80a2e 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
@@ -86,7 +86,6 @@ def homogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
                                        start_vertex_list,
                                        start_vertex_offsets,
                                        h_fan_out,
-                                       num_edge_types,
                                        bool_t with_replacement,
                                        bool_t do_expensive_check,
                                        with_edge_properties=False,
@@ -216,7 +215,7 @@ def homogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
     cdef cugraph_error_t* error_ptr
     cdef uintptr_t ai_fan_out_ptr
 
-    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
index a68e52716e9..1f5ec94705a 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
@@ -83,7 +83,6 @@ def homogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
                                           start_vertex_list,
                                           start_vertex_offsets,
                                           h_fan_out,
-                                          num_edge_types,
                                           bool_t with_replacement,
                                           bool_t do_expensive_check,
                                           with_edge_properties=False,
@@ -210,7 +209,7 @@ def homogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
     cdef cugraph_error_t* error_ptr
     cdef uintptr_t ai_fan_out_ptr
 
-    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create **************
+    # FIXME: refactor the way we are creating pointer. Can use a single helper function to create
 
     assert_CAI_type(start_vertex_list, "start_vertex_list")
     assert_CAI_type(start_vertex_offsets, "start_vertex_offsets", True)

From 95ca286b68fba2c81e98ee0769f96f4212c71116 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 16 Oct 2024 13:45:22 -0700
Subject: [PATCH 069/117] add tests for homogeneous uniform neighborhood
 sampling

---
 ..._homogeneous_uniform_neighbor_sampling.cpp | 423 ++++++++++++++++++
 1 file changed, 423 insertions(+)
 create mode 100644 cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp

diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
new file mode 100644
index 00000000000..b331165dbc2
--- /dev/null
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -0,0 +1,423 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/device_comm_wrapper.hpp"
+#include "utilities/mg_utilities.hpp"
+#include "utilities/property_generator_utilities.hpp"
+#include "utilities/test_graphs.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Homogeneous_Uniform_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  bool with_replacement{true};
+
+  bool edge_masking{false};
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() { handle_ = cugraph::test::initialize_mg_handle(); }
+
+  static void TearDownTestCase() { handle_.reset(); }
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  {
+    auto [homogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
+
+    HighResTimer hr_timer{};
+
+    // 1. create MG graph
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG construct graph");
+    }
+
+    auto [mg_graph, mg_edge_weights, mg_renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, true>(
+        *handle_,
+        input_usecase,
+        true /* test_weighted */,
+        true /* renumber */,
+        false /* drop_self_loops */,
+        false /* drop_multi_edges */);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto mg_graph_view = mg_graph.view();
+    auto mg_edge_weight_view =
+      mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(mg_graph_view), bool>> edge_mask{std::nullopt};
+    if (homogeneous_uniform_neighbor_sampling_usecase.edge_masking) {
+      edge_mask = cugraph::test::generate<decltype(mg_graph_view), bool>::edge_property(
+        *handle_, mg_graph_view, 2);
+      mg_graph_view.attach_edge_mask((*edge_mask).view());
+    }
+
+    //
+    // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
+    // batches
+    //
+
+    constexpr float select_probability{0.05};
+
+    raft::random::RngState rng_state(handle_->get_comms().get_rank());
+
+    auto random_sources = cugraph::select_random_vertices(
+      *handle_,
+      mg_graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      //20,
+      
+      std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
+      
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+    /*
+    rmm::device_uvector<float> random_numbers(random_sources.size(), handle_->get_stream());
+
+    cugraph::detail::uniform_random_fill(handle_->get_stream(),
+                                         random_numbers.data(),
+                                         random_numbers.size(),
+                                         float{0},
+                                         float{1},
+                                         rng_state);
+
+    std::tie(random_numbers, random_sources) = cugraph::test::sort_by_key<float, vertex_t>(
+      *handle_, std::move(random_numbers), std::move(random_sources));
+
+    random_numbers.resize(0, handle_->get_stream());
+    random_numbers.shrink_to_fit(handle_->get_stream());
+    */
+
+    auto seed_sizes = cugraph::host_scalar_allgather(
+      handle_->get_comms(), random_sources.size(), handle_->get_stream());
+    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches = (num_seeds + homogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
+                         homogeneous_uniform_neighbor_sampling_usecase.batch_size;
+
+    std::vector<size_t> seed_offsets(seed_sizes.size());
+    std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
+
+    auto batch_number = cugraph::test::modulo_sequence<int32_t>(
+      *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
+
+    /*
+    rmm::device_uvector<int32_t> unique_batches(num_batches, handle_->get_stream());
+    cugraph::detail::sequence_fill(
+      handle_->get_stream(), unique_batches.data(), unique_batches.size(), int32_t{0});
+    */
+    /*
+    auto comm_ranks = cugraph::test::modulo_sequence<int32_t>(
+      *handle_, num_batches, handle_->get_comms().get_size(), int32_t{0});
+    */
+
+    auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
+      *handle_, random_sources.size(), int32_t{handle_->get_comms().get_rank()});
+    
+    //raft::print_device_vector("random_sources", random_sources.data(), random_sources.size(), std::cout);
+    //raft::print_device_vector("comm_ranks", comm_ranks.data(), comm_ranks.size(), std::cout);
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle_->get_stream());
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::uniform_neighbor_sample(
+        *handle_,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(std::make_tuple(
+          raft::device_span<int32_t const>{unique_batches.data(), unique_batches.size()},
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()})),
+        raft::host_span<int32_t const>(uniform_neighbor_sampling_usecase.fanout.data(),
+                                       uniform_neighbor_sampling_usecase.fanout.size()),
+        rng_state,
+        true,
+        uniform_neighbor_sampling_usecase.with_replacement),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG uniform_neighbor_sample");
+    }
+    RAFT_CUDA_TRY(cudaDeviceSynchronize());
+    std::cout << "running homogeneous uniform neighbor sample" << std::endl;
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::homogeneous_uniform_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_uniform_neighbor_sampling_usecase.with_replacement
+        }
+        );
+    std::cout << "Done running homogeneous uniform neighbor sample" << std::endl;
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (homogeneous_uniform_neighbor_sampling_usecase.check_correctness) {
+      // Consolidate results on GPU 0
+      auto mg_start_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()});
+      auto mg_aggregate_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_out.data(), src_out.size()});
+      auto mg_aggregate_dst = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_out.data(), dst_out.size()});
+      auto mg_aggregate_wgt =
+        wgt_out ? std::make_optional(cugraph::test::device_gatherv(
+                    *handle_, raft::device_span<weight_t const>{wgt_out->data(), wgt_out->size()}))
+                : std::nullopt;
+
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(
+        vertices.data(), mg_aggregate_src.data(), mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(vertices.data() + mg_aggregate_src.size(),
+                 mg_aggregate_dst.data(),
+                 mg_aggregate_dst.size(),
+                 handle_->get_stream());
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      vertices = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        *handle_, std::move(vertices), mg_graph_view.vertex_partition_range_lasts());
+
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle_->get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle_->get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle_->get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle_->get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        *handle_,
+        mg_graph_view,
+        mg_edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      auto mg_aggregate_src_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_compare.data(), src_compare.size()});
+      auto mg_aggregate_dst_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_compare.data(), dst_compare.size()});
+      auto mg_aggregate_wgt_compare =
+        wgt_compare
+          ? std::make_optional(cugraph::test::device_gatherv(
+              *handle_,
+              raft::device_span<weight_t const>{wgt_compare->data(), wgt_compare->size()}))
+          : std::nullopt;
+
+      if (handle_->get_comms().get_rank() == 0) {
+        cugraph::test::validate_extracted_graph_is_subgraph(*handle_,
+                                                            mg_aggregate_src_compare,
+                                                            mg_aggregate_dst_compare,
+                                                            mg_aggregate_wgt_compare,
+                                                            mg_aggregate_src,
+                                                            mg_aggregate_dst,
+                                                            mg_aggregate_wgt);
+
+        if (random_sources.size() < 100) {
+          // This validation is too expensive for large number of vertices
+          if (mg_aggregate_src.size() > 0) {
+            cugraph::test::validate_sampling_depth(*handle_,
+                                                   std::move(mg_aggregate_src),
+                                                   std::move(mg_aggregate_dst),
+                                                   std::move(mg_aggregate_wgt),
+                                                   std::move(mg_start_src),
+                                                   homogeneous_uniform_neighbor_sampling_usecase.fanout.size());
+          }
+        }
+      }
+    }
+#endif
+  }
+
+ private:
+  static std::unique_ptr<raft::handle_t> handle_;
+};
+
+template <typename input_usecase_t>
+std::unique_ptr<raft::handle_t> Tests_MGHomogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ =
+  nullptr;
+
+using Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File =
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+//#if 0
+using Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat =
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+//#endif
+
+#if 0
+TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+#endif
+
+//#if 0
+TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(
+      // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
+      cugraph::test::Rmat_Usecase(5, 16, 0.57, 0.19, 0.19, 0, false, false))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
+                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
+//#endif
+
+CUGRAPH_MG_TEST_PROGRAM_MAIN()

From 383bfc402ea2fec6b3932d333af4bc8863863379 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 15:05:18 -0700
Subject: [PATCH 070/117] add method to fill a buffer array with a scalar

---
 cpp/tests/utilities/thrust_wrapper.cu  | 22 ++++++++++++++++++++++
 cpp/tests/utilities/thrust_wrapper.hpp |  5 +++++
 2 files changed, 27 insertions(+)

diff --git a/cpp/tests/utilities/thrust_wrapper.cu b/cpp/tests/utilities/thrust_wrapper.cu
index ef1c4f831eb..2d759a22b1f 100644
--- a/cpp/tests/utilities/thrust_wrapper.cu
+++ b/cpp/tests/utilities/thrust_wrapper.cu
@@ -41,6 +41,7 @@ namespace test {
 
 template <typename value_t>
 cugraph::dataframe_buffer_type_t<value_t> sort(
+  
   raft::handle_t const& handle, cugraph::dataframe_buffer_type_t<value_t> const& values)
 {
   auto sorted_values = cugraph::allocate_dataframe_buffer<value_t>(
@@ -403,6 +404,27 @@ template rmm::device_uvector<int64_t> sequence(raft::handle_t const& handle,
                                                size_t repeat_count,
                                                int64_t init);
 
+template <typename value_t>
+cugraph::dataframe_buffer_type_t<value_t> scalar_fill(raft::handle_t const& handle,
+                                                      size_t length,
+                                                      value_t value)
+{
+  auto values = cugraph::allocate_dataframe_buffer<value_t>(length, handle.get_stream());
+  
+  thrust::tabulate(handle.get_thrust_policy(),
+                    values.begin(),
+                    values.end(),
+                    [value] __device__(size_t i) {
+                      return value;
+                    });
+  
+  return values;
+}
+
+template rmm::device_uvector<int32_t> scalar_fill(raft::handle_t const& handle,
+                                                  size_t length,
+                                                  int32_t value);
+
 template <typename value_t>
 cugraph::dataframe_buffer_type_t<value_t> modulo_sequence(raft::handle_t const& handle,
                                                           size_t length,
diff --git a/cpp/tests/utilities/thrust_wrapper.hpp b/cpp/tests/utilities/thrust_wrapper.hpp
index afdff33d80a..3d23d92d6cb 100644
--- a/cpp/tests/utilities/thrust_wrapper.hpp
+++ b/cpp/tests/utilities/thrust_wrapper.hpp
@@ -73,6 +73,11 @@ cugraph::dataframe_buffer_type_t<value_t> sequence(raft::handle_t const& handle,
                                                    size_t repeat_count,
                                                    value_t init);
 
+template <typename value_t>
+cugraph::dataframe_buffer_type_t<value_t> scalar_fill(raft::handle_t const& handle,
+                                                      size_t length,
+                                                      value_t value);
+
 // return (init + i) % modulo, where i = [0, length)
 template <typename value_t>
 cugraph::dataframe_buffer_type_t<value_t> modulo_sequence(raft::handle_t const& handle,

From 68fa2f193c41baae92aae99d39a5a55f90d3220b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 15:06:24 -0700
Subject: [PATCH 071/117] add method to sort and count unique elements in
 buffer array

---
 cpp/src/detail/utility_wrappers_32.cu    |  4 ++++
 cpp/src/detail/utility_wrappers_64.cu    |  4 ++++
 cpp/src/detail/utility_wrappers_impl.cuh | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/cpp/src/detail/utility_wrappers_32.cu b/cpp/src/detail/utility_wrappers_32.cu
index bd5927c8cb3..67e86b21213 100644
--- a/cpp/src/detail/utility_wrappers_32.cu
+++ b/cpp/src/detail/utility_wrappers_32.cu
@@ -63,6 +63,10 @@ template void scalar_fill(raft::handle_t const& handle, size_t* d_value, size_t
 
 template void scalar_fill(raft::handle_t const& handle, float* d_value, size_t size, float value);
 
+template void sort(raft::handle_t const& handle, int32_t* d_value, size_t size);
+
+template size_t unique(raft::handle_t const& handle, int32_t* d_value, size_t size);
+
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int32_t* d_value,
                             size_t size,
diff --git a/cpp/src/detail/utility_wrappers_64.cu b/cpp/src/detail/utility_wrappers_64.cu
index 31f021a7ff5..e75fdf7b075 100644
--- a/cpp/src/detail/utility_wrappers_64.cu
+++ b/cpp/src/detail/utility_wrappers_64.cu
@@ -61,6 +61,10 @@ template void scalar_fill(raft::handle_t const& handle,
 
 template void scalar_fill(raft::handle_t const& handle, double* d_value, size_t size, double value);
 
+template void sort(raft::handle_t const& handle, int64_t* d_value, size_t size);
+
+template size_t unique(raft::handle_t const& handle, int32_t* d_value, size_t size);
+
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int64_t* d_value,
                             size_t size,
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index 9f6e5cbb135..e28039a3c02 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -36,6 +36,7 @@
 #include <thrust/transform.h>
 #include <thrust/transform_reduce.h>
 #include <thrust/tuple.h>
+#include <thrust/unique.h>
 
 namespace cugraph {
 namespace detail {
@@ -63,6 +64,23 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
   thrust::fill_n(handle.get_thrust_policy(), d_value, size, value);
 }
 
+template <typename value_t>
+void sort(raft::handle_t const& handle, value_t* d_value, size_t size)
+{
+  thrust::sort(handle.get_thrust_policy(), d_value, d_value + size);
+}
+
+template <typename value_t>
+size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size)
+{
+  //auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
+  auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
+  //auto num_unique_element = 
+  return thrust::distance(d_value, unique_element_last);
+  //masked_edgelist_srcs.resize(2* masked_edgelist_srcs.size(), handle.get_stream());
+}
+
+
 template <typename value_t>
 void sequence_fill(rmm::cuda_stream_view const& stream_view,
                    value_t* d_value,

From 18899ed587ff3b43aff53213243892cd65c3f865 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 15:07:00 -0700
Subject: [PATCH 072/117] add method to sort and count unique elements in
 buffer array

---
 .../cugraph/detail/utility_wrappers.hpp       | 29 +++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index fa61c91a4ab..76b362e58bf 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -65,6 +65,35 @@ void uniform_random_fill(rmm::cuda_stream_view const& stream_view,
 template <typename value_t>
 void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, value_t value);
 
+/**
+ * @brief    Sort a buffer
+ *
+ * @tparam      value_t      type of the value to operate on
+ *
+ * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
+ * and handles to various CUDA libraries) to run graph algorithms.
+ * @param[out]  d_value      device array to sort
+ * @param[in]   size         number of elements in array
+ *
+ */
+template <typename value_t>
+void sort(raft::handle_t const& handle, value_t* d_value, size_t size);
+
+/**
+ * @brief    Keep unique element from a buffer
+ *
+ * @tparam      value_t      type of the value to operate on
+ *
+ * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
+ * and handles to various CUDA libraries) to run graph algorithms.
+ * @param[in]  d_value      device array of unique elements.
+ * @param[in]   size         number of elements in array
+ * @return the number of unique elements.
+ *
+ */
+template <typename value_t>
+size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size);
+
 
 /**
  * @brief    Increment the values of a buffer by a constant value

From 57e6f969f9ac9c1eee99972a2016285d8d262388 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 15:08:38 -0700
Subject: [PATCH 073/117] update computation of map from label to comm rank

---
 cpp/src/c_api/neighbor_sampling.cpp | 40 +++++++++++++++++++++++------
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 39ecec73f96..d7b4990fea2 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -893,17 +893,41 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
       if constexpr (multi_gpu) {
         if (start_vertex_labels) {
-
-          rmm::device_uvector<label_t> label_to_comm_rank_d_vector((*start_vertex_labels).size(), handle_.get_stream());
+          rmm::device_uvector<label_t> unique_labels((*start_vertex_labels).size(), handle_.get_stream());
+          raft::copy(unique_labels.data(),
+                     (*start_vertex_labels).data(),
+                     unique_labels.size(),
+                     handle_.get_stream());
+
+          // Get unique labels
+          // sort the start_vertex_labels
+          cugraph::detail::sort(handle_.get_stream(), unique_labels.begin(), unique_labels.size());
+          auto num_unique_labels = cugraph::detail::unique(
+            handle_.get_stream(), unique_labels.begin(), unique_labels.size());
+          
+          (*label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 
           cugraph::detail::scalar_fill(handle_.get_stream(),
-                                       label_to_comm_rank_d_vector.begin(),
-                                       label_to_comm_rank_d_vector.size(),
+                                       (*label_to_comm_rank).begin(), // This should be rename to rank
+                                       (*label_to_comm_rank).size(),
                                        label_t{handle_.get_comms().get_rank()});
           
-          std::tie(start_vertices, *start_vertex_labels, *label_to_comm_rank) =
-            cugraph::detail::shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-              handle_, std::move(start_vertices), std::move(*start_vertex_labels), std::move(label_to_comm_rank_d_vector));
+          // Perform allgather to get global_label_to_comm_rank_d_vector (not ere, after shuffling)
+          
+          // comm_rank instead of label_to_comm_rank
+          // THe C++ code is not expecting that I fill that with each vertex. need to fill that in by where the label goes
+          // size of the unique vertex labels
+
+          //Solution: Take the global start vertex labels on each GPU, 
+          // Perform sort unique, to get unoqiue labels,
+          // Perform allgather to get the same [0, 1, 2, 3], need 3 zeros and 3 1s
+          // The most important thing is to get the right lenght on each machine. Perform a thrust::fill (rank),
+          // then an allgatherv to get [0, 0, 0, 1, 1], [0, 0, 0, 1, 1]
+
+
+          std::tie(start_vertices, *start_vertex_labels) =
+            cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
+              handle_, std::move(start_vertices), std::move(*start_vertex_labels));
           
         } else {
           start_vertices =
@@ -2082,7 +2106,7 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets, // RENAME?
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,

From d34d85c944a58ab90c2735f40c7f87a0349d9007 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 21:14:04 -0700
Subject: [PATCH 074/117] perform allgatherv of the local mapping from label to
 comm rank

---
 cpp/src/c_api/neighbor_sampling.cpp | 36 ++++++++++++++++-------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index d7b4990fea2..7d535f45058 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -866,7 +866,8 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                  handle_.get_stream());
 
       std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
-      std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{std::nullopt};
+      std::optional<rmm::device_uvector<label_t>> local_label_to_comm_rank{std::nullopt};
+      std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{std::nullopt}; // global after allgatherv
 
       if (start_vertex_offsets_ != nullptr) {
         // Get the number of labels on each GPU
@@ -905,26 +906,29 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           auto num_unique_labels = cugraph::detail::unique(
             handle_.get_stream(), unique_labels.begin(), unique_labels.size());
           
-          (*label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
+          (*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 
           cugraph::detail::scalar_fill(handle_.get_stream(),
-                                       (*label_to_comm_rank).begin(), // This should be rename to rank
-                                       (*label_to_comm_rank).size(),
+                                       (*local_label_to_comm_rank).begin(), // This should be rename to rank
+                                       (*local_label_to_comm_rank).size(),
                                        label_t{handle_.get_comms().get_rank()});
           
-          // Perform allgather to get global_label_to_comm_rank_d_vector (not ere, after shuffling)
+          // Perform allgather to get global_label_to_comm_rank_d_vector
+          auto recvcounts = cugraph::host_scalar_allgather(
+            handle_.get_comms(), num_unique_labels, handle_.get_stream());
           
-          // comm_rank instead of label_to_comm_rank
-          // THe C++ code is not expecting that I fill that with each vertex. need to fill that in by where the label goes
-          // size of the unique vertex labels
-
-          //Solution: Take the global start vertex labels on each GPU, 
-          // Perform sort unique, to get unoqiue labels,
-          // Perform allgather to get the same [0, 1, 2, 3], need 3 zeros and 3 1s
-          // The most important thing is to get the right lenght on each machine. Perform a thrust::fill (rank),
-          // then an allgatherv to get [0, 0, 0, 1, 1], [0, 0, 0, 1, 1]
-
-
+          std::vector<size_t> displacements(recvcounts.size());
+           std::exclusive_scan(recvcounts.begin(), recvcounts.end(), displacements.begin(), size_t{0});
+          
+          (*label_to_comm_rank).resize(displacements.back() + recvcounts.back(), handle_.get_stream());
+
+          cugraph::device_allgatherv(handle_.get_comms(),
+                                   (*local_label_to_comm_rank).begin(),
+                                   (*label_to_comm_rank).begin(),
+                                   recvcounts,
+                                   displacements,
+                                   handle_.get_stream());
+        
           std::tie(start_vertices, *start_vertex_labels) =
             cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
               handle_, std::move(start_vertices), std::move(*start_vertex_labels));

From 2aa0903526d8b4c9779192915fa2336172f1e9a8 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 21:42:59 -0700
Subject: [PATCH 075/117] udpate tests for
 'mg_homogeneous_uniform_neighbor_sampling'

---
 ..._homogeneous_uniform_neighbor_sampling.cpp | 42 +++++++++----------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
index b331165dbc2..178aaaf1825 100644
--- a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -144,6 +144,21 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
 
     auto batch_number = cugraph::test::modulo_sequence<int32_t>(
       *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
+    
+    // Get unique batch_number -> label_list
+    rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
+
+    raft::copy(label_list.data(),
+               batch_number.data(),
+               batch_number.size(),
+               handle_->get_stream());
+    
+    label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
+    label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
+
+    auto num_unique_labels = label_list.size();
+    
+    // 
 
     /*
     rmm::device_uvector<int32_t> unique_batches(num_batches, handle_->get_stream());
@@ -156,7 +171,12 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
     */
 
     auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
-      *handle_, random_sources.size(), int32_t{handle_->get_comms().get_rank()});
+      *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
+
+    // perform allgatherv
+    comm_ranks =
+      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+    
     
     //raft::print_device_vector("random_sources", random_sources.data(), random_sources.size(), std::cout);
     //raft::print_device_vector("comm_ranks", comm_ranks.data(), comm_ranks.size(), std::cout);
@@ -325,28 +345,9 @@ std::unique_ptr<raft::handle_t> Tests_MGHomogeneous_Uniform_Neighbor_Sampling<in
 using Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File =
   Tests_MGHomogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
 
-//#if 0
 using Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat =
   Tests_MGHomogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
 
-//#endif
-
-#if 0
-TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
-{
-  run_current_test<int32_t, int32_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
-INSTANTIATE_TEST_SUITE_P(
-  file_test,
-  Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File,
-  ::testing::Combine(
-    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true}),
-    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
-#endif
-
-//#if 0
 TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
   run_current_test<int32_t, int32_t, float>(
@@ -418,6 +419,5 @@ INSTANTIATE_TEST_SUITE_P(
                       Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
                       Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
-//#endif
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()

From fa0cb8858ff2107a483ebdb2dca8f4722ba21e47 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 22:26:22 -0700
Subject: [PATCH 076/117] update neighbor sampling call for 'NO_CUGRAPH_OPS'

---
 .../heterogeneous_biased_neighbor_sampling.cpp | 18 ++++++++++++++----
 ...heterogeneous_uniform_neighbor_sampling.cpp | 17 +++++++++++++----
 .../homogeneous_biased_neighbor_sampling.cpp   | 12 +++++++++---
 .../homogeneous_uniform_neighbor_sampling.cpp  | 11 ++++++++---
 4 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index e3fd0b841fb..745a9049c93 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -130,10 +130,15 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
     EXPECT_THROW(
       cugraph::heterogeneous_biased_neighbor_sample(
         handle,
+        rng_state,
         graph_view,
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        *edge_weight_view,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
@@ -141,9 +146,14 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        rng_state,
-        true,
-        heterogeneous_biased_neighbor_sampling_usecase.flag_replacement),
+        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_biased_neighbor_sampling_usecase.flag_replacement
+        }
+        ),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 53c346226f2..3a5d64b5ded 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -130,10 +130,14 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
     EXPECT_THROW(
       cugraph::heterogeneous_uniform_neighbor_sample(
         handle,
+        rng_state,
         graph_view,
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
@@ -141,9 +145,14 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        rng_state,
-        true,
-        heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement),
+        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement
+        }
+        ),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
index 61d418e2f4d..a3f6b612337 100644
--- a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -121,10 +121,12 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
     EXPECT_THROW(
       cugraph::homogeneous_biased_neighbor_sample(
         handle,
+        rng_state,
         graph_view,
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        *edge_weight_view,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
@@ -132,9 +134,13 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        rng_state,
-        true,
-        homogeneous_biased_neighbor_sampling_usecase.flag_replacement),
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_biased_neighbor_sampling_usecase.flag_replacement
+        }
+        ),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
index 789ac48111a..5b18ecf6a50 100644
--- a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -121,6 +121,7 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
     EXPECT_THROW(
       cugraph::homogeneous_uniform_neighbor_sample(
         handle,
+        rng_state,
         graph_view,
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
@@ -132,9 +133,13 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        rng_state,
-        true,
-        homogeneous_uniform_neighbor_sampling_usecase.flag_replacement),
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_uniform_neighbor_sampling_usecase.flag_replacement
+        }
+        ),
       std::exception);
 #else
     if (cugraph::test::g_perf) {

From 990d2edf0acadc3e633f4d8d811a11f162a293cc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 22:48:24 -0700
Subject: [PATCH 077/117] remove outdated code

---
 ..._homogeneous_uniform_neighbor_sampling.cpp | 60 +++++--------------
 1 file changed, 16 insertions(+), 44 deletions(-)

diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
index 178aaaf1825..7132231cf15 100644
--- a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -116,22 +116,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
     //
     //  Now we'll assign the vertices to batches
     //
-    /*
-    rmm::device_uvector<float> random_numbers(random_sources.size(), handle_->get_stream());
-
-    cugraph::detail::uniform_random_fill(handle_->get_stream(),
-                                         random_numbers.data(),
-                                         random_numbers.size(),
-                                         float{0},
-                                         float{1},
-                                         rng_state);
-
-    std::tie(random_numbers, random_sources) = cugraph::test::sort_by_key<float, vertex_t>(
-      *handle_, std::move(random_numbers), std::move(random_sources));
-
-    random_numbers.resize(0, handle_->get_stream());
-    random_numbers.shrink_to_fit(handle_->get_stream());
-    */
 
     auto seed_sizes = cugraph::host_scalar_allgather(
       handle_->get_comms(), random_sources.size(), handle_->get_stream());
@@ -157,18 +141,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
     label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
 
     auto num_unique_labels = label_list.size();
-    
-    // 
-
-    /*
-    rmm::device_uvector<int32_t> unique_batches(num_batches, handle_->get_stream());
-    cugraph::detail::sequence_fill(
-      handle_->get_stream(), unique_batches.data(), unique_batches.size(), int32_t{0});
-    */
-    /*
-    auto comm_ranks = cugraph::test::modulo_sequence<int32_t>(
-      *handle_, num_batches, handle_->get_comms().get_size(), int32_t{0});
-    */
 
     auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
       *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
@@ -176,10 +148,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
     // perform allgatherv
     comm_ranks =
       cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
-    
-    
-    //raft::print_device_vector("random_sources", random_sources.data(), random_sources.size(), std::cout);
-    //raft::print_device_vector("comm_ranks", comm_ranks.data(), comm_ranks.size(), std::cout);
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
 
@@ -190,23 +158,28 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
 
 #ifdef NO_CUGRAPH_OPS
     EXPECT_THROW(
-      cugraph::uniform_neighbor_sample(
+      cugraph::homogeneous_uniform_neighbor_sample(
         *handle_,
+        rng_state,
         mg_graph_view,
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(std::make_tuple(
-          raft::device_span<int32_t const>{unique_batches.data(), unique_batches.size()},
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()})),
-        raft::host_span<int32_t const>(uniform_neighbor_sampling_usecase.fanout.data(),
-                                       uniform_neighbor_sampling_usecase.fanout.size()),
-        rng_state,
-        true,
-        uniform_neighbor_sampling_usecase.with_replacement),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_uniform_neighbor_sampling_usecase.with_replacement
+        }
+        ),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -215,7 +188,7 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
       hr_timer.start("MG uniform_neighbor_sample");
     }
     RAFT_CUDA_TRY(cudaDeviceSynchronize());
-    std::cout << "running homogeneous uniform neighbor sample" << std::endl;
+  
     auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
       cugraph::homogeneous_uniform_neighbor_sample(
         *handle_,
@@ -239,7 +212,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
             homogeneous_uniform_neighbor_sampling_usecase.with_replacement
         }
         );
-    std::cout << "Done running homogeneous uniform neighbor sample" << std::endl;
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement

From d44a46f7955a32741dd4c40eff9ff456b63b20ac Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Mon, 21 Oct 2024 22:49:21 -0700
Subject: [PATCH 078/117] udpate tests for
 'mg_homogeneous_biased_neighbor_sampling'

---
 ...g_homogeneous_biased_neighbor_sampling.cpp | 396 ++++++++++++++++++
 1 file changed, 396 insertions(+)
 create mode 100644 cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp

diff --git a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
new file mode 100644
index 00000000000..1ac6c4244ce
--- /dev/null
+++ b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
@@ -0,0 +1,396 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/device_comm_wrapper.hpp"
+#include "utilities/mg_utilities.hpp"
+#include "utilities/property_generator_utilities.hpp"
+#include "utilities/test_graphs.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Homogeneous_Biased_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  bool with_replacement{true};
+
+  bool edge_masking{false};
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_MGHomogeneous_Biased_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() { handle_ = cugraph::test::initialize_mg_handle(); }
+
+  static void TearDownTestCase() { handle_.reset(); }
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  {
+    auto [homogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
+
+    HighResTimer hr_timer{};
+
+    // 1. create MG graph
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG construct graph");
+    }
+
+    auto [mg_graph, mg_edge_weights, mg_renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, true>(
+        *handle_,
+        input_usecase,
+        true /* test_weighted */,
+        true /* renumber */,
+        false /* drop_self_loops */,
+        false /* drop_multi_edges */);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto mg_graph_view = mg_graph.view();
+    auto mg_edge_weight_view =
+      mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
+
+    std::optional<cugraph::edge_property_t<decltype(mg_graph_view), bool>> edge_mask{std::nullopt};
+    if (homogeneous_biased_neighbor_sampling_usecase.edge_masking) {
+      edge_mask = cugraph::test::generate<decltype(mg_graph_view), bool>::edge_property(
+        *handle_, mg_graph_view, 2);
+      mg_graph_view.attach_edge_mask((*edge_mask).view());
+    }
+
+    //
+    // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
+    // batches
+    //
+
+    constexpr float select_probability{0.05};
+
+    raft::random::RngState rng_state(handle_->get_comms().get_rank());
+
+    auto random_sources = cugraph::select_random_vertices(
+      *handle_,
+      mg_graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      //20,
+      
+      std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
+      
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto seed_sizes = cugraph::host_scalar_allgather(
+      handle_->get_comms(), random_sources.size(), handle_->get_stream());
+    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches = (num_seeds + homogeneous_biased_neighbor_sampling_usecase.batch_size - 1) /
+                         homogeneous_biased_neighbor_sampling_usecase.batch_size;
+
+    std::vector<size_t> seed_offsets(seed_sizes.size());
+    std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
+
+    auto batch_number = cugraph::test::modulo_sequence<int32_t>(
+      *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
+    
+    // Get unique batch_number -> label_list
+    rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
+
+    raft::copy(label_list.data(),
+               batch_number.data(),
+               batch_number.size(),
+               handle_->get_stream());
+    
+    label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
+    label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
+
+    auto num_unique_labels = label_list.size();
+
+    auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
+      *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
+
+    // perform allgatherv
+    comm_ranks =
+      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle_->get_stream());
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      cugraph::homogeneous_biased_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        mg_edge_weight_view
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_biased_neighbor_sampling_usecase.with_replacement
+        }
+        ),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG homogeneous_biased_neighbor_sample");
+    }
+    RAFT_CUDA_TRY(cudaDeviceSynchronize());
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::homogeneous_biased_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
+        *mg_edge_weight_view,
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            homogeneous_biased_neighbor_sampling_usecase.with_replacement
+        }
+        );
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (homogeneous_biased_neighbor_sampling_usecase.check_correctness) {
+      // Consolidate results on GPU 0
+      auto mg_start_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()});
+      auto mg_aggregate_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_out.data(), src_out.size()});
+      auto mg_aggregate_dst = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_out.data(), dst_out.size()});
+      auto mg_aggregate_wgt =
+        wgt_out ? std::make_optional(cugraph::test::device_gatherv(
+                    *handle_, raft::device_span<weight_t const>{wgt_out->data(), wgt_out->size()}))
+                : std::nullopt;
+
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(
+        vertices.data(), mg_aggregate_src.data(), mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(vertices.data() + mg_aggregate_src.size(),
+                 mg_aggregate_dst.data(),
+                 mg_aggregate_dst.size(),
+                 handle_->get_stream());
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      vertices = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        *handle_, std::move(vertices), mg_graph_view.vertex_partition_range_lasts());
+
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle_->get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle_->get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle_->get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle_->get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        *handle_,
+        mg_graph_view,
+        mg_edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      auto mg_aggregate_src_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_compare.data(), src_compare.size()});
+      auto mg_aggregate_dst_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_compare.data(), dst_compare.size()});
+      auto mg_aggregate_wgt_compare =
+        wgt_compare
+          ? std::make_optional(cugraph::test::device_gatherv(
+              *handle_,
+              raft::device_span<weight_t const>{wgt_compare->data(), wgt_compare->size()}))
+          : std::nullopt;
+
+      if (handle_->get_comms().get_rank() == 0) {
+        cugraph::test::validate_extracted_graph_is_subgraph(*handle_,
+                                                            mg_aggregate_src_compare,
+                                                            mg_aggregate_dst_compare,
+                                                            mg_aggregate_wgt_compare,
+                                                            mg_aggregate_src,
+                                                            mg_aggregate_dst,
+                                                            mg_aggregate_wgt);
+
+        if (random_sources.size() < 100) {
+          // This validation is too expensive for large number of vertices
+          if (mg_aggregate_src.size() > 0) {
+            cugraph::test::validate_sampling_depth(*handle_,
+                                                   std::move(mg_aggregate_src),
+                                                   std::move(mg_aggregate_dst),
+                                                   std::move(mg_aggregate_wgt),
+                                                   std::move(mg_start_src),
+                                                   homogeneous_biased_neighbor_sampling_usecase.fanout.size());
+          }
+        }
+      }
+    }
+#endif
+  }
+
+ private:
+  static std::unique_ptr<raft::handle_t> handle_;
+};
+
+template <typename input_usecase_t>
+std::unique_ptr<raft::handle_t> Tests_MGHomogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ =
+  nullptr;
+
+using Tests_MGHomogeneous_Biased_Neighbor_Sampling_File =
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+using Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat =
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true}),
+    ::testing::Values(
+      // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
+      cugraph::test::Rmat_Usecase(5, 16, 0.57, 0.19, 0.19, 0, false, false))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
+                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
+
+CUGRAPH_MG_TEST_PROGRAM_MAIN()

From a1a6180d4ad58b68bb3ac464c64484e8d3ac15ec Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 08:39:14 -0700
Subject: [PATCH 079/117] add mg tests for heterogeneous uniform and biased
 neighborhood sampling

---
 ...heterogeneous_biased_neighbor_sampling.cpp | 411 ++++++++++++++++++
 ...eterogeneous_uniform_neighbor_sampling.cpp | 410 +++++++++++++++++
 2 files changed, 821 insertions(+)
 create mode 100644 cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
 create mode 100644 cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp

diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
new file mode 100644
index 00000000000..d9cdee3479c
--- /dev/null
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -0,0 +1,411 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/device_comm_wrapper.hpp"
+#include "utilities/mg_utilities.hpp"
+#include "utilities/property_generator_utilities.hpp"
+#include "utilities/test_graphs.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  int32_t num_edge_types{1};
+  bool with_replacement{true};
+
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() { handle_ = cugraph::test::initialize_mg_handle(); }
+
+  static void TearDownTestCase() { handle_.reset(); }
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  {
+    using edge_type_t = int32_t;
+
+    auto [heterogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
+
+    HighResTimer hr_timer{};
+
+    // 1. create MG graph
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG construct graph");
+    }
+
+    auto [mg_graph, mg_edge_weights, mg_renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, true>(
+        *handle_,
+        input_usecase,
+        true /* test_weighted */,
+        true /* renumber */,
+        false /* drop_self_loops */,
+        false /* drop_multi_edges */);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto mg_graph_view = mg_graph.view();
+    auto mg_edge_weight_view =
+      mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
+
+
+    //
+    // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
+    // batches
+    //
+
+    constexpr float select_probability{0.05};
+
+    raft::random::RngState rng_state(handle_->get_comms().get_rank());
+
+    auto random_sources = cugraph::select_random_vertices(
+      *handle_,
+      mg_graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      //20,
+      
+      std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
+      
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto seed_sizes = cugraph::host_scalar_allgather(
+      handle_->get_comms(), random_sources.size(), handle_->get_stream());
+    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches = (num_seeds + heterogeneous_biased_neighbor_sampling_usecase.batch_size - 1) /
+                         heterogeneous_biased_neighbor_sampling_usecase.batch_size;
+
+    std::vector<size_t> seed_offsets(seed_sizes.size());
+    std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
+
+    auto batch_number = cugraph::test::modulo_sequence<int32_t>(
+      *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
+    
+    // Get unique batch_number -> label_list
+    rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
+
+    raft::copy(label_list.data(),
+               batch_number.data(),
+               batch_number.size(),
+               handle_->get_stream());
+    
+    label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
+    label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
+
+    auto num_unique_labels = label_list.size();
+
+    auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
+      *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
+
+    // perform allgatherv
+    comm_ranks =
+      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle_->get_stream());
+    
+    // Generate the edge types
+
+    std::optional<cugraph::edge_property_t<decltype(mg_graph_view), edge_type_t>> edge_types{
+      std::nullopt};
+
+    if (heterogeneous_biased_neighbor_sampling_usecase.num_edge_types > 1) {
+      edge_types = cugraph::test::generate<decltype(mg_graph_view), edge_type_t>::edge_property(
+          *handle_,
+          mg_graph_view,
+          heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
+    }
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      ccugraph::heterogeneous_biased_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_biased_neighbor_sampling_usecase.with_replacement
+        }
+        ),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG biased_neighbor_sample");
+    }
+    RAFT_CUDA_TRY(cudaDeviceSynchronize());
+  
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::heterogeneous_biased_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        *mg_edge_weight_view,
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_biased_neighbor_sampling_usecase.with_replacement
+        }
+        );
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (heterogeneous_biased_neighbor_sampling_usecase.check_correctness) {
+      // Consolidate results on GPU 0
+      auto mg_start_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()});
+      auto mg_aggregate_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_out.data(), src_out.size()});
+      auto mg_aggregate_dst = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_out.data(), dst_out.size()});
+      auto mg_aggregate_wgt =
+        wgt_out ? std::make_optional(cugraph::test::device_gatherv(
+                    *handle_, raft::device_span<weight_t const>{wgt_out->data(), wgt_out->size()}))
+                : std::nullopt;
+
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(
+        vertices.data(), mg_aggregate_src.data(), mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(vertices.data() + mg_aggregate_src.size(),
+                 mg_aggregate_dst.data(),
+                 mg_aggregate_dst.size(),
+                 handle_->get_stream());
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      vertices = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        *handle_, std::move(vertices), mg_graph_view.vertex_partition_range_lasts());
+
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle_->get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle_->get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle_->get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle_->get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        *handle_,
+        mg_graph_view,
+        mg_edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      auto mg_aggregate_src_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_compare.data(), src_compare.size()});
+      auto mg_aggregate_dst_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_compare.data(), dst_compare.size()});
+      auto mg_aggregate_wgt_compare =
+        wgt_compare
+          ? std::make_optional(cugraph::test::device_gatherv(
+              *handle_,
+              raft::device_span<weight_t const>{wgt_compare->data(), wgt_compare->size()}))
+          : std::nullopt;
+
+      if (handle_->get_comms().get_rank() == 0) {
+        cugraph::test::validate_extracted_graph_is_subgraph(*handle_,
+                                                            mg_aggregate_src_compare,
+                                                            mg_aggregate_dst_compare,
+                                                            mg_aggregate_wgt_compare,
+                                                            mg_aggregate_src,
+                                                            mg_aggregate_dst,
+                                                            mg_aggregate_wgt);
+
+        if (random_sources.size() < 100) {
+          // This validation is too expensive for large number of vertices
+          if (mg_aggregate_src.size() > 0) {
+            cugraph::test::validate_sampling_depth(*handle_,
+                                                   std::move(mg_aggregate_src),
+                                                   std::move(mg_aggregate_dst),
+                                                   std::move(mg_aggregate_wgt),
+                                                   std::move(mg_start_src),
+                                                   heterogeneous_biased_neighbor_sampling_usecase.fanout.size());
+          }
+        }
+      }
+    }
+#endif
+  }
+
+ private:
+  static std::unique_ptr<raft::handle_t> handle_;
+};
+
+template <typename input_usecase_t>
+std::unique_ptr<raft::handle_t> Tests_MGHeterogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ =
+  nullptr;
+
+using Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File =
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+using Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat =
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(
+      // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
+      cugraph::test::Rmat_Usecase(5, 16, 0.57, 0.19, 0.19, 0, false, false))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
+
+CUGRAPH_MG_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
new file mode 100644
index 00000000000..87dab4801fb
--- /dev/null
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -0,0 +1,410 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "detail/nbr_sampling_validate.hpp"
+#include "utilities/base_fixture.hpp"
+#include "utilities/device_comm_wrapper.hpp"
+#include "utilities/mg_utilities.hpp"
+#include "utilities/property_generator_utilities.hpp"
+#include "utilities/test_graphs.hpp"
+
+#include <cugraph/sampling_functions.hpp>
+#include <cugraph/utilities/high_res_timer.hpp>
+
+#include <gtest/gtest.h>
+
+struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
+  std::vector<int32_t> fanout{{-1}};
+  int32_t batch_size{10};
+  int32_t num_edge_types{1};
+  bool with_replacement{true};
+
+  bool check_correctness{true};
+};
+
+template <typename input_usecase_t>
+class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
+  : public ::testing::TestWithParam<
+      std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t>> {
+ public:
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling() {}
+
+  static void SetUpTestCase() { handle_ = cugraph::test::initialize_mg_handle(); }
+
+  static void TearDownTestCase() { handle_.reset(); }
+
+  virtual void SetUp() {}
+  virtual void TearDown() {}
+
+  template <typename vertex_t, typename edge_t, typename weight_t>
+  void run_current_test(std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  {
+    using edge_type_t = int32_t;
+
+    auto [heterogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
+
+    HighResTimer hr_timer{};
+
+    // 1. create MG graph
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG construct graph");
+    }
+
+    auto [mg_graph, mg_edge_weights, mg_renumber_map_labels] =
+      cugraph::test::construct_graph<vertex_t, edge_t, weight_t, false, true>(
+        *handle_,
+        input_usecase,
+        true /* test_weighted */,
+        true /* renumber */,
+        false /* drop_self_loops */,
+        false /* drop_multi_edges */);
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    auto mg_graph_view = mg_graph.view();
+    auto mg_edge_weight_view =
+      mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
+
+
+    //
+    // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
+    // batches
+    //
+
+    constexpr float select_probability{0.05};
+
+    raft::random::RngState rng_state(handle_->get_comms().get_rank());
+
+    auto random_sources = cugraph::select_random_vertices(
+      *handle_,
+      mg_graph_view,
+      std::optional<raft::device_span<vertex_t const>>{std::nullopt},
+      rng_state,
+      //20,
+      
+      std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
+               std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
+      
+      false,
+      false);
+
+    //
+    //  Now we'll assign the vertices to batches
+    //
+
+    auto seed_sizes = cugraph::host_scalar_allgather(
+      handle_->get_comms(), random_sources.size(), handle_->get_stream());
+    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches = (num_seeds + heterogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
+                         heterogeneous_uniform_neighbor_sampling_usecase.batch_size;
+
+    std::vector<size_t> seed_offsets(seed_sizes.size());
+    std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
+
+    auto batch_number = cugraph::test::modulo_sequence<int32_t>(
+      *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
+    
+    // Get unique batch_number -> label_list
+    rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
+
+    raft::copy(label_list.data(),
+               batch_number.data(),
+               batch_number.size(),
+               handle_->get_stream());
+    
+    label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
+    label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
+
+    auto num_unique_labels = label_list.size();
+
+    auto comm_ranks = cugraph::test::scalar_fill<int32_t>(
+      *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
+
+    // perform allgatherv
+    comm_ranks =
+      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+
+    rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
+
+    raft::copy(random_sources_copy.data(),
+               random_sources.data(),
+               random_sources.size(),
+               handle_->get_stream());
+    
+    // Generate the edge types
+
+    std::optional<cugraph::edge_property_t<decltype(mg_graph_view), edge_type_t>> edge_types{
+      std::nullopt};
+
+    if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
+      edge_types = cugraph::test::generate<decltype(mg_graph_view), edge_type_t>::edge_property(
+          *handle_,
+          mg_graph_view,
+          heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
+    }
+
+#ifdef NO_CUGRAPH_OPS
+    EXPECT_THROW(
+      ccugraph::heterogeneous_uniform_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_uniform_neighbor_sampling_usecase.with_replacement
+        }
+        ),
+      std::exception);
+#else
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.start("MG uniform_neighbor_sample");
+    }
+    RAFT_CUDA_TRY(cudaDeviceSynchronize());
+  
+    auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
+      cugraph::heterogeneous_uniform_neighbor_sample(
+        *handle_,
+        rng_state,
+        mg_graph_view,
+        mg_edge_weight_view,
+        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
+        edge_types
+            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                        .view()}
+            : std::nullopt,
+        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        std::make_optional(
+          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
+        std::make_optional(
+          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
+        cugraph::sampling_flags_t{
+            cugraph::prior_sources_behavior_t{0},
+            true, // return_hops
+            false, // dedupe_sources
+            heterogeneous_uniform_neighbor_sampling_usecase.with_replacement
+        }
+        );
+
+    if (cugraph::test::g_perf) {
+      RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
+      handle_->get_comms().barrier();
+      hr_timer.stop();
+      hr_timer.display_and_clear(std::cout);
+    }
+
+    if (heterogeneous_uniform_neighbor_sampling_usecase.check_correctness) {
+      // Consolidate results on GPU 0
+      auto mg_start_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()});
+      auto mg_aggregate_src = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_out.data(), src_out.size()});
+      auto mg_aggregate_dst = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_out.data(), dst_out.size()});
+      auto mg_aggregate_wgt =
+        wgt_out ? std::make_optional(cugraph::test::device_gatherv(
+                    *handle_, raft::device_span<weight_t const>{wgt_out->data(), wgt_out->size()}))
+                : std::nullopt;
+
+      //  First validate that the extracted edges are actually a subset of the
+      //  edges in the input graph
+      rmm::device_uvector<vertex_t> vertices(2 * mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(
+        vertices.data(), mg_aggregate_src.data(), mg_aggregate_src.size(), handle_->get_stream());
+      raft::copy(vertices.data() + mg_aggregate_src.size(),
+                 mg_aggregate_dst.data(),
+                 mg_aggregate_dst.size(),
+                 handle_->get_stream());
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      vertices = cugraph::detail::shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
+        *handle_, std::move(vertices), mg_graph_view.vertex_partition_range_lasts());
+
+      vertices = cugraph::test::sort<vertex_t>(*handle_, std::move(vertices));
+      vertices = cugraph::test::unique<vertex_t>(*handle_, std::move(vertices));
+
+      rmm::device_uvector<size_t> d_subgraph_offsets(2, handle_->get_stream());
+      std::vector<size_t> h_subgraph_offsets({0, vertices.size()});
+
+      raft::update_device(d_subgraph_offsets.data(),
+                          h_subgraph_offsets.data(),
+                          h_subgraph_offsets.size(),
+                          handle_->get_stream());
+
+      rmm::device_uvector<vertex_t> src_compare(0, handle_->get_stream());
+      rmm::device_uvector<vertex_t> dst_compare(0, handle_->get_stream());
+      std::optional<rmm::device_uvector<weight_t>> wgt_compare{std::nullopt};
+      std::tie(src_compare, dst_compare, wgt_compare, std::ignore) = extract_induced_subgraphs(
+        *handle_,
+        mg_graph_view,
+        mg_edge_weight_view,
+        raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
+        raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
+        true);
+
+      auto mg_aggregate_src_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{src_compare.data(), src_compare.size()});
+      auto mg_aggregate_dst_compare = cugraph::test::device_gatherv(
+        *handle_, raft::device_span<vertex_t const>{dst_compare.data(), dst_compare.size()});
+      auto mg_aggregate_wgt_compare =
+        wgt_compare
+          ? std::make_optional(cugraph::test::device_gatherv(
+              *handle_,
+              raft::device_span<weight_t const>{wgt_compare->data(), wgt_compare->size()}))
+          : std::nullopt;
+
+      if (handle_->get_comms().get_rank() == 0) {
+        cugraph::test::validate_extracted_graph_is_subgraph(*handle_,
+                                                            mg_aggregate_src_compare,
+                                                            mg_aggregate_dst_compare,
+                                                            mg_aggregate_wgt_compare,
+                                                            mg_aggregate_src,
+                                                            mg_aggregate_dst,
+                                                            mg_aggregate_wgt);
+
+        if (random_sources.size() < 100) {
+          // This validation is too expensive for large number of vertices
+          if (mg_aggregate_src.size() > 0) {
+            cugraph::test::validate_sampling_depth(*handle_,
+                                                   std::move(mg_aggregate_src),
+                                                   std::move(mg_aggregate_dst),
+                                                   std::move(mg_aggregate_wgt),
+                                                   std::move(mg_start_src),
+                                                   heterogeneous_uniform_neighbor_sampling_usecase.fanout.size());
+          }
+        }
+      }
+    }
+#endif
+  }
+
+ private:
+  static std::unique_ptr<raft::handle_t> handle_;
+};
+
+template <typename input_usecase_t>
+std::unique_ptr<raft::handle_t> Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ =
+  nullptr;
+
+using Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File =
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
+
+using Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat =
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
+
+TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_File_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
+{
+  run_current_test<int32_t, int32_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
+{
+  run_current_test<int32_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
+{
+  run_current_test<int64_t, int64_t, float>(
+    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  file_test,
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  file_large_test,
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_small_test,
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
+    ::testing::Values(
+      // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
+      cugraph::test::Rmat_Usecase(5, 16, 0.57, 0.19, 0.19, 0, false, false))));
+
+INSTANTIATE_TEST_SUITE_P(
+  rmat_benchmark_test, /* note that scale & edge factor can be overridden in benchmarking (with
+                          --gtest_filter to select only the rmat_benchmark_test with a specific
+                          vertex & edge type combination) by command line arguments and do not
+                          include more than one Rmat_Usecase that differ only in scale or edge
+                          factor (to avoid running same benchmarks more than once) */
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat,
+  ::testing::Combine(
+    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+    ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
+
+CUGRAPH_MG_TEST_PROGRAM_MAIN()

From 36ce4fcf20e38923b707c843ff486c4808bcc152 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 08:52:24 -0700
Subject: [PATCH 080/117] add new tests to CMakeLists

---
 cpp/tests/CMakeLists.txt | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index bfcab9a026b..3145a2ea1be 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -767,6 +767,26 @@ if(BUILD_CUGRAPH_MG_TESTS)
     # - MG UNIFORM NBR SAMPLING tests -------------------------------------------------------------
     ConfigureTestMG(MG_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/mg_uniform_neighbor_sampling.cpp)
 
+    ###############################################################################################
+    # - MG HOMOGENEOUS UNIFORM NBR SAMPLING tests -------------------------------------------------
+    ConfigureTestMG(
+        MG_HOMOGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/mg_homogeneous_uniform_neighbor_sampling.cpp)
+    
+    ###############################################################################################
+    # - MG HETEROGENEOUS UNIFORM NBR SAMPLING tests -------------------------------------------------
+    ConfigureTestMG(
+        MG_HETEROGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp)
+    
+    ###############################################################################################
+    # - MG HOMOGENEOUS BIASED NBR SAMPLING tests --------------------------------------------------
+    ConfigureTestMG(
+        MG_HOMOGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/mg_homogeneous_biased_neighbor_sampling.cpp)
+    
+    ###############################################################################################
+    # - MG HETEROGENEOUS BIASED NBR SAMPLING tests --------------------------------------------------
+    ConfigureTestMG(
+        MG_HETEROGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/mg_heterogeneous_biased_neighbor_sampling.cpp)
+
     ###############################################################################################
     # - MG BIASED NBR SAMPLING tests --------------------------------------------------------------
     ConfigureTestMG(MG_BIASED_NEIGHBOR_SAMPLING_TEST sampling/mg_biased_neighbor_sampling.cpp)

From c0a618fd9bf78c7c9cfb0d46be89de3528b6ebc9 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 08:55:05 -0700
Subject: [PATCH 081/117] remove unsued variable

---
 cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp | 2 --
 1 file changed, 2 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 3a5d64b5ded..d94d7ad7e5a 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -75,8 +75,6 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
     auto edge_weight_view =
       edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
 
-    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
-
     constexpr float select_probability{0.05};
 
     // FIXME:  Update the tests to initialize RngState and use it instead

From 965001ba3e6a6a363cba0732c230bf9b2043a709 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 09:37:26 -0700
Subject: [PATCH 082/117] fix illegal memory access

---
 cpp/src/sampling/detail/conversion_utilities.cu       | 2 +-
 cpp/src/sampling/detail/conversion_utilities_impl.cuh | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index bc6cb128561..a067005514d 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -41,7 +41,7 @@ rmm::device_uvector<int32_t> flatten_label_map(
                                      label_t{0},
                                      thrust::maximum<label_t>());
 
-  label_map.resize(max_label, handle.get_stream());
+  label_map.resize(max_label + 1, handle.get_stream());
 
   thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
   thrust::scatter(handle.get_thrust_policy(),
diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
index 3ed11d05181..0c8d8ac95ea 100644
--- a/cpp/src/sampling/detail/conversion_utilities_impl.cuh
+++ b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
@@ -42,10 +42,10 @@ rmm::device_uvector<int32_t> flatten_label_map(
                                      label_t{0},
                                      thrust::maximum<label_t>());
 
-  label_map.resize(max_label, handle.get_stream());
+  label_map.resize(max_label + 1, handle.get_stream());
 
-  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), label_t{0});
-  thrust::gather(handle.get_thrust_policy(),
+  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
+  thrust::scatter(handle.get_thrust_policy(),
                  std::get<0>(label_to_output_comm_rank).begin(),
                  std::get<0>(label_to_output_comm_rank).end(),
                  std::get<1>(label_to_output_comm_rank).begin(),

From 11f9c4020547431c4d985cf4684354edf85252d7 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 09:44:06 -0700
Subject: [PATCH 083/117] fix style

---
 .../cugraph/detail/shuffle_wrappers.hpp       |  13 +-
 .../cugraph/detail/utility_wrappers.hpp       |   1 -
 cpp/include/cugraph/sampling_functions.hpp    |   6 +-
 cpp/include/cugraph_c/sampling_algorithms.h   |  59 ++-
 cpp/src/c_api/array.hpp                       |   2 +-
 cpp/src/c_api/neighbor_sampling.cpp           | 432 +++++++++---------
 cpp/src/detail/utility_wrappers_impl.cuh      |   8 +-
 .../detail/conversion_utilities_impl.cuh      |   8 +-
 cpp/src/sampling/neighbor_sampling_impl.hpp   | 232 +++++-----
 .../sampling/neighbor_sampling_mg_v32_e32.cu  |  51 +--
 .../sampling/neighbor_sampling_sg_v32_e32.cu  |  48 +-
 cpp/src/utilities/shuffle_vertices.cuh        |  75 +--
 .../utilities/shuffle_vertices_mg_v32_fp.cu   |  26 +-
 .../shuffle_vertices_mg_v32_integral.cu       |  29 +-
 .../utilities/shuffle_vertices_mg_v64_fp.cu   |  26 +-
 .../shuffle_vertices_mg_v64_integral.cu       |  28 +-
 cpp/tests/CMakeLists.txt                      |   6 +-
 ...heterogeneous_biased_neighbor_sampling.cpp |  99 ++--
 ...eterogeneous_uniform_neighbor_sampling.cpp |  98 ++--
 .../homogeneous_biased_neighbor_sampling.cpp  |  60 ++-
 .../homogeneous_uniform_neighbor_sampling.cpp |  60 ++-
 ...heterogeneous_biased_neighbor_sampling.cpp | 106 ++---
 ...eterogeneous_uniform_neighbor_sampling.cpp | 107 ++---
 ...g_homogeneous_biased_neighbor_sampling.cpp |  83 ++--
 ..._homogeneous_uniform_neighbor_sampling.cpp |  88 ++--
 cpp/tests/utilities/thrust_wrapper.cu         |  16 +-
 .../dask/sampling/uniform_neighbor_sample.py  |   2 +-
 .../heterogeneous_biased_neighbor_sample.py   |  37 +-
 .../heterogeneous_uniform_neighbor_sample.py  |  37 +-
 ...ogeneous_biased_uniform_neighbor_sample.py |  35 +-
 .../homogeneous_uniform_neighbor_sample.py    |  35 +-
 python/pylibcugraph/pylibcugraph/__init__.py  |  22 +-
 .../heterogeneous_biased_neighbor_sample.pyx  |   8 +-
 .../heterogeneous_uniform_neighbor_sample.pyx |   6 +-
 .../homogeneous_biased_neighbor_sample.pyx    |   6 +-
 .../homogeneous_uniform_neighbor_sample.pyx   |   4 +-
 36 files changed, 949 insertions(+), 1010 deletions(-)

diff --git a/cpp/include/cugraph/detail/shuffle_wrappers.hpp b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
index 68ebdf01296..6592150e748 100644
--- a/cpp/include/cugraph/detail/shuffle_wrappers.hpp
+++ b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
@@ -158,12 +158,13 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
  * @return Tuple of vectors storing shuffled vertex & value pairs.
  */
 template <typename vertex_t, typename value0_t, typename value1_t>
-std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<vertex_t>&& vertices,
-  rmm::device_uvector<value0_t>&& values_0,
-  rmm::device_uvector<value1_t>&& values_1);
+std::
+  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<vertex_t>&& vertices,
+    rmm::device_uvector<value0_t>&& values_0,
+    rmm::device_uvector<value1_t>&& values_1);
 
 /**
  * @brief Permute a range.
diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index 76b362e58bf..a7400a43039 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -94,7 +94,6 @@ void sort(raft::handle_t const& handle, value_t* d_value, size_t size);
 template <typename value_t>
 size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size);
 
-
 /**
  * @brief    Increment the values of a buffer by a constant value
  *
diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index de636f84ff1..f3f9b126513 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -317,7 +317,8 @@ struct sampling_flags_t {
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types + edge_type_id]'
+ * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types +
+ * edge_type_id]'
  * @param num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
  * sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param flags A set of flags indicating which sampling features should be used.
@@ -405,7 +406,8 @@ heterogeneous_uniform_neighbor_sample(
  * @param label_to_output_comm_rank Optional device span identifying which rank should get each
  * vertex label.  This should be the same on each rank.
  * @param fan_out Host span defining branching out (fan-out) degree per source vertex for each
- * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types + edge_type_id]'
+ * level. The fanout value at hop x is given by the expression 'fanout[x*num_edge_types +
+ * edge_type_id]'
  * @param num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
  * sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param flags A set of flags indicating which sampling features should be used.
diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index abc8f1c4082..fb0b93a0bed 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -331,10 +331,9 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
  * @brief     Opaque neighborhood sampling heterogeneous fan_out type
  */
 
-
 /**
  * @brief     Uniform Neighborhood Sampling
- * 
+ *
  * @deprecated  This API will be deleted, use cugraph_homogeneous_uniform_neighbor_sample
  *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
@@ -362,8 +361,8 @@ void cugraph_sampling_options_free(cugraph_sampling_options_t* options);
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fan_out values of type INT32
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
  * @param [in,out] rng_state State of the random number generator, updated with each call
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
@@ -391,9 +390,9 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
 
 /**
  * @brief     Biased Neighborhood Sampling
- * 
+ *
  * @deprecated  This API will be deleted, use cugraph_homogeneous_biased_neighbor_sample.
- * 
+ *
  * Returns a sample of the neighborhood around specified start vertices.  Optionally, each
  * start vertex can be associated with a label, allowing the caller to specify multiple batches
  * of sampling requests in the same function call - which should improve GPU utilization.
@@ -422,8 +421,8 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
  * label_to_comm_rank[i].  If not specified then the output data will not be shuffled between ranks.
  * @param [in]  label_offsets Device array of the offsets for each label in the seed list.  This
  *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                           We only support fan_out values of type INT32
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
  * @param [in,out] rng_state State of the random number generator, updated with each call
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
@@ -467,13 +466,12 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
  *                           needs to be transposed
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                            We only support fan_out values of type INT32
- * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
- *                            sample whereas a value greater than 1 translates to heterogeneous neighbor
- *                            sample.
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.
+ * This parameter is only used with the retain_seeds option.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
+ * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -496,7 +494,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
-
 /**
  * @brief     Heterogeneous Biased Neighborhood Sampling
  *
@@ -516,13 +513,12 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
  * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                            We only support fan_out values of type INT32
- * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous neighbor
- *                            sample whereas a value greater than 1 translates to heterogeneous neighbor
- *                            sample.
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.
+ * This parameter is only used with the retain_seeds option.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
+ * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -563,10 +559,10 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
  * @param [in]  graph        Pointer to graph.  NOTE: Graph might be modified if the storage
  *                           needs to be transposed
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                            We only support fan_out values of type INT32
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.
+ * This parameter is only used with the retain_seeds option.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -607,10 +603,10 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
  * @param [in]  edge_biases  Device array of edge biases to use for sampling.  If NULL
  * use the edge weight as the bias. If set to NULL, edges will be sampled uniformly.
  * @param [in]  start_vertices Device array of start vertices for the sampling
- * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.  This
- *                            parameter is only used with the retain_seeds option.
- * @param [in]  fan_out       Host array defining the fan out at each step in the sampling algorithm.
- *                            We only support fan_out values of type INT32
+ * @param [in]  start_vertex_offsets Device array of the offsets for each label in the seed list.
+ * This parameter is only used with the retain_seeds option.
+ * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
+ * algorithm. We only support fan_out values of type INT32
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -633,7 +629,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
-
 /**
  * @deprecated This call should be replaced with cugraph_sample_result_get_majors
  * @brief     Get the source vertices from the sampling algorithm result
diff --git a/cpp/src/c_api/array.hpp b/cpp/src/c_api/array.hpp
index 2fa58d917ae..0ab30a1cb72 100644
--- a/cpp/src/c_api/array.hpp
+++ b/cpp/src/c_api/array.hpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 7d535f45058..5b05b542015 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-
 #include "c_api/abstract_functor.hpp"
 #include "c_api/graph.hpp"
+#include "c_api/graph_helper.hpp"
 #include "c_api/properties.hpp"
 #include "c_api/random.hpp"
 #include "c_api/resource_handle.hpp"
 #include "c_api/utils.hpp"
-#include "c_api/graph_helper.hpp"
+#include "sampling/detail/sampling_utils.hpp"
 
 #include <cugraph_c/algorithms.h>
 #include <cugraph_c/sampling_algorithms.h>
@@ -30,7 +30,6 @@
 #include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
 #include <cugraph/sampling_functions.hpp>
-#include "sampling/detail/sampling_utils.hpp"
 
 #include <raft/core/handle.hpp>
 
@@ -780,18 +779,17 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
   bool do_expensive_check_{false};
   cugraph::c_api::cugraph_sample_result_t* result_{nullptr};
 
-  neighbor_sampling_functor(
-    cugraph_resource_handle_t const* handle,
-    cugraph_rng_state_t* rng_state,
-    cugraph_graph_t* graph,
-    cugraph_edge_property_view_t const* edge_biases,
-    cugraph_type_erased_device_array_view_t const* start_vertices,
-    cugraph_type_erased_device_array_view_t const* start_vertex_offsets,
-    cugraph_type_erased_host_array_view_t const* fan_out,
-    int num_edge_types,
-    cugraph::c_api::cugraph_sampling_options_t options,
-    bool is_biased,
-    bool do_expensive_check)
+  neighbor_sampling_functor(cugraph_resource_handle_t const* handle,
+                            cugraph_rng_state_t* rng_state,
+                            cugraph_graph_t* graph,
+                            cugraph_edge_property_view_t const* edge_biases,
+                            cugraph_type_erased_device_array_view_t const* start_vertices,
+                            cugraph_type_erased_device_array_view_t const* start_vertex_offsets,
+                            cugraph_type_erased_host_array_view_t const* fan_out,
+                            int num_edge_types,
+                            cugraph::c_api::cugraph_sampling_options_t options,
+                            bool is_biased,
+                            bool do_expensive_check)
     : abstract_functor(),
       handle_(*reinterpret_cast<cugraph::c_api::cugraph_resource_handle_t const*>(handle)->handle_),
       rng_state_(reinterpret_cast<cugraph::c_api::cugraph_rng_state_t*>(rng_state)),
@@ -867,17 +865,18 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
       std::optional<rmm::device_uvector<label_t>> start_vertex_labels{std::nullopt};
       std::optional<rmm::device_uvector<label_t>> local_label_to_comm_rank{std::nullopt};
-      std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{std::nullopt}; // global after allgatherv
+      std::optional<rmm::device_uvector<label_t>> label_to_comm_rank{
+        std::nullopt};  // global after allgatherv
 
       if (start_vertex_offsets_ != nullptr) {
         // Get the number of labels on each GPU
         auto num_local_labels = start_vertex_offsets_->size_ - 1;
-        auto global_labels = cugraph::host_scalar_allgather(
+        auto global_labels    = cugraph::host_scalar_allgather(
           handle_.get_comms(), num_local_labels, handle_.get_stream());
-        
+
         std::exclusive_scan(
-            global_labels.begin(), global_labels.end(), global_labels.begin(), size_t{0});
-        
+          global_labels.begin(), global_labels.end(), global_labels.begin(), size_t{0});
+
         // Compute the global start_vertex_label_offsets
         cugraph::detail::transform_increment(handle_.get_stream(),
                                              start_vertex_offsets_->as_type<size_t>(),
@@ -885,16 +884,16 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                                              global_labels[handle_.get_comms().get_rank()]);
 
         // Retrieve the start_vertex_labels
-        start_vertex_labels =
-          cugraph::detail::convert_starting_vertex_offsets_to_labels(
-            handle_,
-            raft::device_span<size_t const>{start_vertex_offsets_->as_type<size_t>(),
-                                            start_vertex_offsets_->size_});
+        start_vertex_labels = cugraph::detail::convert_starting_vertex_offsets_to_labels(
+          handle_,
+          raft::device_span<size_t const>{start_vertex_offsets_->as_type<size_t>(),
+                                          start_vertex_offsets_->size_});
       }
 
       if constexpr (multi_gpu) {
         if (start_vertex_labels) {
-          rmm::device_uvector<label_t> unique_labels((*start_vertex_labels).size(), handle_.get_stream());
+          rmm::device_uvector<label_t> unique_labels((*start_vertex_labels).size(),
+                                                     handle_.get_stream());
           raft::copy(unique_labels.data(),
                      (*start_vertex_labels).data(),
                      unique_labels.size(),
@@ -905,34 +904,37 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           cugraph::detail::sort(handle_.get_stream(), unique_labels.begin(), unique_labels.size());
           auto num_unique_labels = cugraph::detail::unique(
             handle_.get_stream(), unique_labels.begin(), unique_labels.size());
-          
+
           (*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 
-          cugraph::detail::scalar_fill(handle_.get_stream(),
-                                       (*local_label_to_comm_rank).begin(), // This should be rename to rank
-                                       (*local_label_to_comm_rank).size(),
-                                       label_t{handle_.get_comms().get_rank()});
-          
+          cugraph::detail::scalar_fill(
+            handle_.get_stream(),
+            (*local_label_to_comm_rank).begin(),  // This should be rename to rank
+            (*local_label_to_comm_rank).size(),
+            label_t{handle_.get_comms().get_rank()});
+
           // Perform allgather to get global_label_to_comm_rank_d_vector
           auto recvcounts = cugraph::host_scalar_allgather(
             handle_.get_comms(), num_unique_labels, handle_.get_stream());
-          
+
           std::vector<size_t> displacements(recvcounts.size());
-           std::exclusive_scan(recvcounts.begin(), recvcounts.end(), displacements.begin(), size_t{0});
-          
-          (*label_to_comm_rank).resize(displacements.back() + recvcounts.back(), handle_.get_stream());
+          std::exclusive_scan(
+            recvcounts.begin(), recvcounts.end(), displacements.begin(), size_t{0});
+
+          (*label_to_comm_rank)
+            .resize(displacements.back() + recvcounts.back(), handle_.get_stream());
 
           cugraph::device_allgatherv(handle_.get_comms(),
-                                   (*local_label_to_comm_rank).begin(),
-                                   (*label_to_comm_rank).begin(),
-                                   recvcounts,
-                                   displacements,
-                                   handle_.get_stream());
-        
+                                     (*local_label_to_comm_rank).begin(),
+                                     (*label_to_comm_rank).begin(),
+                                     recvcounts,
+                                     displacements,
+                                     handle_.get_stream());
+
           std::tie(start_vertices, *start_vertex_labels) =
             cugraph::detail::shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
               handle_, std::move(start_vertices), std::move(*start_vertex_labels));
-          
+
         } else {
           start_vertices =
             cugraph::detail::shuffle_ext_vertices_to_local_gpu_by_vertex_partitioning(
@@ -951,7 +953,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         graph_view.local_vertex_partition_range_first(),
         graph_view.local_vertex_partition_range_last(),
         do_expensive_check_);
-      
+
       rmm::device_uvector<vertex_t> src(0, handle_.get_stream());
       rmm::device_uvector<vertex_t> dst(0, handle_.get_stream());
       std::optional<rmm::device_uvector<weight_t>> wgt{std::nullopt};
@@ -960,7 +962,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       std::optional<rmm::device_uvector<int32_t>> hop{std::nullopt};
       std::optional<rmm::device_uvector<label_t>> edge_label{std::nullopt};
       std::optional<rmm::device_uvector<size_t>> offsets{std::nullopt};
-  
+
       // FIXME: For biased sampling, the user should pass either biases or edge weights,
       // otherwised throw an error and suggest the user to call uniform neighbor sample instead
 
@@ -979,21 +981,17 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
                 ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
-                                                                      (*start_vertex_labels).size())
-                : std::nullopt,
-              label_to_comm_rank
-                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
-                                                                  (*label_to_comm_rank).size()})
+                                                                   (*start_vertex_labels).size())
                 : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              label_to_comm_rank ? std::make_optional(raft::device_span<int const>{
+                                     (*label_to_comm_rank).data(), (*label_to_comm_rank).size()})
+                                 : std::nullopt,
+              raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
               num_edge_types_,
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
+              cugraph::sampling_flags_t{options_.prior_sources_behavior_,
+                                        options_.return_hops_,
+                                        options_.dedupe_sources_,
+                                        options_.with_replacement_},
               do_expensive_check_);
         } else {
           auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
@@ -1007,23 +1005,18 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
               raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
               (start_vertex_offsets_ != nullptr)
                 ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
-                                                                      (*start_vertex_labels).size())
-                : std::nullopt,
-              label_to_comm_rank
-                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
-                                                                  (*label_to_comm_rank).size()})
+                                                                   (*start_vertex_labels).size())
                 : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
+              label_to_comm_rank ? std::make_optional(raft::device_span<int const>{
+                                     (*label_to_comm_rank).data(), (*label_to_comm_rank).size()})
+                                 : std::nullopt,
+              raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
               num_edge_types_,
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
+              cugraph::sampling_flags_t{options_.prior_sources_behavior_,
+                                        options_.return_hops_,
+                                        options_.dedupe_sources_,
+                                        options_.with_replacement_},
               do_expensive_check_);
-
         }
       } else {
         // Call homogeneous neighbor sample
@@ -1042,47 +1035,39 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                 ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
                                                                    (*start_vertex_labels).size())
                 : std::nullopt,
-              label_to_comm_rank
-                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
-                                                                  (*label_to_comm_rank).size()})
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
+              label_to_comm_rank ? std::make_optional(raft::device_span<int const>{
+                                     (*label_to_comm_rank).data(), (*label_to_comm_rank).size()})
+                                 : std::nullopt,
+              raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
+              cugraph::sampling_flags_t{options_.prior_sources_behavior_,
+                                        options_.return_hops_,
+                                        options_.dedupe_sources_,
+                                        options_.with_replacement_},
               do_expensive_check_);
         } else {
           if (is_biased_) {
-          auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
-            cugraph::homogeneous_uniform_neighbor_sample(
-              handle_,
-              rng_state_->rng_state_,
-              graph_view,
-              (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
-              (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
-              (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
-              raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
-              (start_vertex_offsets_ != nullptr)
-                ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
-                                                                      (*start_vertex_labels).size())
-                : std::nullopt,
-              label_to_comm_rank
-                ? std::make_optional(raft::device_span<int const>{(*label_to_comm_rank).data(),
-                                                                  (*label_to_comm_rank).size()})
-                : std::nullopt,
-              raft::host_span<const int>(
-                                        fan_out_->as_type<const int>(), fan_out_->size_),
-              cugraph::sampling_flags_t{
-                options_.prior_sources_behavior_,
-                options_.return_hops_,
-                options_.dedupe_sources_,
-                options_.with_replacement_
-              },
-              do_expensive_check_);
+            auto&& [src, dst, wgt, edge_id, edge_type, hop, offsets] =
+              cugraph::homogeneous_uniform_neighbor_sample(
+                handle_,
+                rng_state_->rng_state_,
+                graph_view,
+                (edge_weights != nullptr) ? std::make_optional(edge_weights->view()) : std::nullopt,
+                (edge_ids != nullptr) ? std::make_optional(edge_ids->view()) : std::nullopt,
+                (edge_types != nullptr) ? std::make_optional(edge_types->view()) : std::nullopt,
+                raft::device_span<vertex_t const>{start_vertices.data(), start_vertices.size()},
+                (start_vertex_offsets_ != nullptr)
+                  ? std::make_optional<raft::device_span<int const>>((*start_vertex_labels).data(),
+                                                                     (*start_vertex_labels).size())
+                  : std::nullopt,
+                label_to_comm_rank ? std::make_optional(raft::device_span<int const>{
+                                       (*label_to_comm_rank).data(), (*label_to_comm_rank).size()})
+                                   : std::nullopt,
+                raft::host_span<const int>(fan_out_->as_type<const int>(), fan_out_->size_),
+                cugraph::sampling_flags_t{options_.prior_sources_behavior_,
+                                          options_.return_hops_,
+                                          options_.dedupe_sources_,
+                                          options_.with_replacement_},
+                do_expensive_check_);
           }
         }
       }
@@ -1125,20 +1110,20 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
       }
 
       if (options_.renumber_results_) {
-        if (num_edge_types_ == 1) { // homogeneous renumbering
+        if (num_edge_types_ == 1) {  // homogeneous renumbering
           if (options_.compression_type_ == cugraph_compression_type_t::COO) {
             // COO
 
             rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
             rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
             std::tie(output_majors,
-                    minors,
-                    wgt,
-                    edge_id,
-                    edge_type,
-                    label_hop_offsets,
-                    output_renumber_map,
-                    renumber_map_offsets) =
+                     minors,
+                     wgt,
+                     edge_id,
+                     edge_type,
+                     label_hop_offsets,
+                     output_renumber_map,
+                     renumber_map_offsets) =
               cugraph::renumber_and_sort_sampled_edgelist<vertex_t>(
                 handle_,
                 std::move(src),
@@ -1151,13 +1136,14 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
-                                          : std::nullopt,
+                options_.retain_seeds_
+                  ? std::make_optional(raft::device_span<size_t const>{
+                      start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
+                  : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
-                offsets ? edge_label->size() : size_t{1}, // FIXME: update edge_label ?
+                offsets ? edge_label->size() : size_t{1},  // FIXME: update edge_label ?
                 hop ? fan_out_->size_ : size_t{1},
                 src_is_major,
                 do_expensive_check_);
@@ -1167,21 +1153,22 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           } else {
             // (D)CSC, (D)CSR
 
-            bool doubly_compress = (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
-                                  (options_.compression_type_ == cugraph_compression_type_t::DCSC);
+            bool doubly_compress =
+              (options_.compression_type_ == cugraph_compression_type_t::DCSR) ||
+              (options_.compression_type_ == cugraph_compression_type_t::DCSC);
 
             rmm::device_uvector<size_t> output_major_offsets(0, handle_.get_stream());
             rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
 
             std::tie(majors,
-                    output_major_offsets,
-                    minors,
-                    wgt,
-                    edge_id,
-                    edge_type,
-                    label_hop_offsets,
-                    output_renumber_map,
-                    renumber_map_offsets) =
+                     output_major_offsets,
+                     minors,
+                     wgt,
+                     edge_id,
+                     edge_type,
+                     label_hop_offsets,
+                     output_renumber_map,
+                     renumber_map_offsets) =
               cugraph::renumber_and_compress_sampled_edgelist<vertex_t>(
                 handle_,
                 std::move(src),
@@ -1194,9 +1181,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                   ? std::make_optional(raft::device_span<vertex_t const>{
                       start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                   : std::nullopt,
-                 options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                          start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
-                                          : std::nullopt,
+                options_.retain_seeds_
+                  ? std::make_optional(raft::device_span<size_t const>{
+                      start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
+                  : std::nullopt,
                 offsets ? std::make_optional(
                             raft::device_span<size_t const>{offsets->data(), offsets->size()})
                         : std::nullopt,
@@ -1214,35 +1202,37 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           // These are now represented by label_hop_offsets
           hop.reset();
           offsets.reset();
-        
-        } else { // heterogeneous renumbering
 
-          rmm::device_uvector<vertex_t> vertex_type_offsets(graph_view.local_vertex_partition_range_size(), handle_.get_stream());
+        } else {  // heterogeneous renumbering
+
+          rmm::device_uvector<vertex_t> vertex_type_offsets(
+            graph_view.local_vertex_partition_range_size(), handle_.get_stream());
 
           cugraph::detail::sequence_fill(handle_.get_stream(),
-                                vertex_type_offsets.begin(),
-                                vertex_type_offsets.size(),
-                                vertex_t{0});
+                                         vertex_type_offsets.begin(),
+                                         vertex_type_offsets.size(),
+                                         vertex_t{0});
 
           rmm::device_uvector<vertex_t> output_majors(0, handle_.get_stream());
           rmm::device_uvector<vertex_t> output_renumber_map(0, handle_.get_stream());
-          
-          std::optional<rmm::device_uvector<edge_t>> renumbered_and_sorted_edge_id_renumber_map(std::nullopt);
-          std::optional<rmm::device_uvector<size_t>>    renumbered_and_sorted_edge_id_renumber_map_label_type_offsets(std::nullopt);
+
+          std::optional<rmm::device_uvector<edge_t>> renumbered_and_sorted_edge_id_renumber_map(
+            std::nullopt);
+          std::optional<rmm::device_uvector<size_t>>
+            renumbered_and_sorted_edge_id_renumber_map_label_type_offsets(std::nullopt);
 
           // extract the edge_type from label_type_hop_offsets
           std::optional<rmm::device_uvector<size_t>> label_type_hop_offsets{std::nullopt};
 
           std::tie(output_majors,
-              minors,
-              wgt,
-              edge_id,
-              label_type_hop_offsets, // Contains information about the type and hop offsets
-              output_renumber_map,
-              (*renumber_map_offsets),
-              renumbered_and_sorted_edge_id_renumber_map,
-              renumbered_and_sorted_edge_id_renumber_map_label_type_offsets
-              ) = 
+                   minors,
+                   wgt,
+                   edge_id,
+                   label_type_hop_offsets,  // Contains information about the type and hop offsets
+                   output_renumber_map,
+                   (*renumber_map_offsets),
+                   renumbered_and_sorted_edge_id_renumber_map,
+                   renumbered_and_sorted_edge_id_renumber_map_label_type_offsets) =
             cugraph::heterogeneous_renumber_and_sort_sampled_edgelist<vertex_t>(
               handle_,
               std::move(src),
@@ -1255,34 +1245,35 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                 ? std::make_optional(raft::device_span<vertex_t const>{
                     start_vertices_->as_type<vertex_t>(), start_vertices_->size_})
                 : std::nullopt,
-              options_.retain_seeds_ ? std::make_optional(raft::device_span<size_t const>{
-                                        start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
-                                    : std::nullopt,
+              options_.retain_seeds_
+                ? std::make_optional(raft::device_span<size_t const>{
+                    start_vertex_offsets_->as_type<size_t>(), start_vertex_offsets_->size_})
+                : std::nullopt,
               offsets ? std::make_optional(
                           raft::device_span<size_t const>{offsets->data(), offsets->size()})
                       : std::nullopt,
-              raft::device_span<vertex_t const>{
-                    vertex_type_offsets.data(), vertex_type_offsets.size()},
-              
+              raft::device_span<vertex_t const>{vertex_type_offsets.data(),
+                                                vertex_type_offsets.size()},
+
               start_vertex_offsets_ ? start_vertex_offsets_->size_ : size_t{1},
               hop ? fan_out_->size_ : size_t{1},
               size_t{1},
               size_t{1},
               src_is_major,
               do_expensive_check_);
-            if (edge_type) {
-              (*edge_type).resize(raft::device_span<size_t const>{
-                (*label_type_hop_offsets).data(), (*label_type_hop_offsets).size()}.back() - 1, handle_.get_stream());
-              cugraph::detail::sequence_fill(handle_.get_stream(),
-                                (*edge_type).begin(),
-                                (*edge_type).size(),
-                                edge_type_t{0});
-            }
-            
+          if (edge_type) {
+            (*edge_type)
+              .resize(raft::device_span<size_t const>{(*label_type_hop_offsets).data(),
+                                                      (*label_type_hop_offsets).size()}
+                          .back() -
+                        1,
+                      handle_.get_stream());
+            cugraph::detail::sequence_fill(
+              handle_.get_stream(), (*edge_type).begin(), (*edge_type).size(), edge_type_t{0});
+          }
 
-            majors.emplace(std::move(output_majors));
-            renumber_map.emplace(std::move(output_renumber_map));
-        
+          majors.emplace(std::move(output_majors));
+          renumber_map.emplace(std::move(output_renumber_map));
         }
 
       } else {
@@ -1290,23 +1281,22 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           CUGRAPH_FAIL("Can only use COO format if not renumbering");
         }
 
-
         std::tie(src, dst, wgt, edge_id, edge_type, label_hop_offsets) =
-          cugraph::sort_sampled_edgelist(handle_,
-                                         std::move(src),
-                                         std::move(dst),
-                                         std::move(wgt),
-                                         std::move(edge_id),
-                                         std::move(edge_type),
-                                         std::move(hop),
-                                         offsets
-                                           ? std::make_optional(raft::device_span<size_t const>{
-                                               offsets->data(), offsets->size()})
-                                           : std::nullopt,
-                                         edge_label ? edge_label->size() : size_t{1}, // FIXME: update edge_label
-                                         hop ? fan_out_->size_ : size_t{1},
-                                         src_is_major,
-                                         do_expensive_check_);
+          cugraph::sort_sampled_edgelist(
+            handle_,
+            std::move(src),
+            std::move(dst),
+            std::move(wgt),
+            std::move(edge_id),
+            std::move(edge_type),
+            std::move(hop),
+            offsets ? std::make_optional(
+                        raft::device_span<size_t const>{offsets->data(), offsets->size()})
+                    : std::nullopt,
+            edge_label ? edge_label->size() : size_t{1},  // FIXME: update edge_label
+            hop ? fan_out_->size_ : size_t{1},
+            src_is_major,
+            do_expensive_check_);
 
         majors.emplace(std::move(src));
         minors = std::move(dst);
@@ -1336,9 +1326,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
         (label_hop_offsets)
           ? new cugraph::c_api::cugraph_type_erased_device_array_t(*label_hop_offsets, SIZE_T)
           : nullptr,
-        (edge_label)
-          ? new cugraph::c_api::cugraph_type_erased_device_array_t(edge_label.value(), INT32) // FIXME: update edge_label
-          : nullptr,
+        (edge_label) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
+                         edge_label.value(), INT32)  // FIXME: update edge_label
+                     : nullptr,
         (renumber_map) ? new cugraph::c_api::cugraph_type_erased_device_array_t(
                            renumber_map.value(), graph_->vertex_type_)
                        : nullptr,
@@ -1918,7 +1908,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_sample_result_t** result,
   cugraph_error_t** error)
 {
-
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
 
   CAPI_EXPECTS(
@@ -1977,7 +1966,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
                                            fan_out,
                                            rng_state,
                                            std::move(options_cpp),
-                                          do_expensive_check};
+                                           do_expensive_check};
   return cugraph::c_api::run_algorithm(graph, functor, result, error);
 }
 
@@ -1995,13 +1984,13 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   cugraph_error_t** error)
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
-  
+
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify start_vertex_offsets if retain_seeds is true",
                *error);
-  
+
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                     start_vertex_offsets)
@@ -2009,7 +1998,7 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-    
+
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
         ->type_ == INT32,
@@ -2061,13 +2050,13 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
     CUGRAPH_INVALID_INPUT,
     "edge_biases is required if the graph is not weighted",
     *error);
-  
+
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify start_vertex_offsets if retain_seeds is true",
                *error);
-  
+
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                     start_vertex_offsets)
@@ -2075,7 +2064,7 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-    
+
   CAPI_EXPECTS(
     reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
         ->type_ == INT32,
@@ -2110,7 +2099,7 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
-  const cugraph_type_erased_device_array_view_t* start_vertex_offsets, // RENAME?
+  const cugraph_type_erased_device_array_view_t* start_vertex_offsets,  // RENAME?
   const cugraph_type_erased_host_array_view_t* fan_out,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
@@ -2118,13 +2107,13 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   cugraph_error_t** error)
 {
   auto options_cpp = *reinterpret_cast<cugraph::c_api::cugraph_sampling_options_t const*>(options);
-  
+
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify start_vertex_offsets if retain_seeds is true",
                *error);
-  
+
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                     start_vertex_offsets)
@@ -2132,14 +2121,14 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-  
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                fan_out)
-                ->type_ == INT32,
-            CUGRAPH_INVALID_INPUT,
-            "fan_out type must be INT32",
-            *error);
-    
+
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out type must be INT32",
+    *error);
+
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                    start_vertices)
@@ -2155,7 +2144,7 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
                                     start_vertices,
                                     start_vertex_offsets,
                                     fan_out,
-                                    1, // num_edge_types
+                                    1,  // num_edge_types
                                     std::move(options_cpp),
                                     FALSE,
                                     do_expensive_check};
@@ -2183,13 +2172,13 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
     CUGRAPH_INVALID_INPUT,
     "edge_biases is required if the graph is not weighted",
     *error);
-  
+
   // FIXME: Should we maintain this contition?
   CAPI_EXPECTS((!options_cpp.retain_seeds_) || (start_vertex_offsets != nullptr),
                CUGRAPH_INVALID_INPUT,
                "must specify start_vertex_offsets if retain_seeds is true",
                *error);
-  
+
   CAPI_EXPECTS((start_vertex_offsets == nullptr) ||
                  (reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                     start_vertex_offsets)
@@ -2197,14 +2186,14 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                CUGRAPH_INVALID_INPUT,
                "start_vertex_offsets should be of type size_t",
                *error);
-  
-  CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(
-                fan_out)
-                ->type_ == INT32,
-            CUGRAPH_INVALID_INPUT,
-            "fan_out type must be INT32",
-            *error);
-    
+
+  CAPI_EXPECTS(
+    reinterpret_cast<cugraph::c_api::cugraph_type_erased_host_array_view_t const*>(fan_out)
+        ->type_ == INT32,
+    CUGRAPH_INVALID_INPUT,
+    "fan_out type must be INT32",
+    *error);
+
   CAPI_EXPECTS(reinterpret_cast<cugraph::c_api::cugraph_graph_t*>(graph)->vertex_type_ ==
                  reinterpret_cast<cugraph::c_api::cugraph_type_erased_device_array_view_t const*>(
                    start_vertices)
@@ -2213,7 +2202,6 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                "vertex type of graph and start_vertices must match",
                *error);
 
-
   neighbor_sampling_functor functor{handle,
                                     rng_state,
                                     graph,
@@ -2221,7 +2209,7 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
                                     start_vertices,
                                     start_vertex_offsets,
                                     fan_out,
-                                    1, // num_edge_types
+                                    1,  // num_edge_types
                                     std::move(options_cpp),
                                     TRUE,
                                     do_expensive_check};
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index e28039a3c02..eaca2516dbf 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -73,14 +73,13 @@ void sort(raft::handle_t const& handle, value_t* d_value, size_t size)
 template <typename value_t>
 size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size)
 {
-  //auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
+  // auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
   auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
-  //auto num_unique_element = 
+  // auto num_unique_element =
   return thrust::distance(d_value, unique_element_last);
-  //masked_edgelist_srcs.resize(2* masked_edgelist_srcs.size(), handle.get_stream());
+  // masked_edgelist_srcs.resize(2* masked_edgelist_srcs.size(), handle.get_stream());
 }
 
-
 template <typename value_t>
 void sequence_fill(rmm::cuda_stream_view const& stream_view,
                    value_t* d_value,
@@ -105,7 +104,6 @@ void transform_increment(rmm::cuda_stream_view const& stream_view,
                     }));
 }
 
-
 template <typename value_t>
 void stride_fill(rmm::cuda_stream_view const& stream_view,
                  value_t* d_value,
diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
index 0c8d8ac95ea..6eff0154581 100644
--- a/cpp/src/sampling/detail/conversion_utilities_impl.cuh
+++ b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
@@ -46,10 +46,10 @@ rmm::device_uvector<int32_t> flatten_label_map(
 
   thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
   thrust::scatter(handle.get_thrust_policy(),
-                 std::get<0>(label_to_output_comm_rank).begin(),
-                 std::get<0>(label_to_output_comm_rank).end(),
-                 std::get<1>(label_to_output_comm_rank).begin(),
-                 label_map.begin());
+                  std::get<0>(label_to_output_comm_rank).begin(),
+                  std::get<0>(label_to_output_comm_rank).end(),
+                  std::get<1>(label_to_output_comm_rank).begin(),
+                  label_map.begin());
 
   return label_map;
 }
diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index 5f72a4d99a8..e3a06f323d1 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -18,12 +18,10 @@
 
 #include "prims/fill_edge_property.cuh"
 #include "prims/transform_e.cuh"
-
 #include "sampling/detail/sampling_utils.hpp"
 
 #include <cugraph/detail/shuffle_wrappers.hpp>
 #include <cugraph/detail/utility_wrappers.hpp>
-
 #include <cugraph/graph.hpp>
 #include <cugraph/graph_functions.hpp>
 #include <cugraph/sampling_functions.hpp>
@@ -101,16 +99,19 @@ neighbor_sample_impl(raft::handle_t const& handle,
     "Invalid input argument: number of levels should not overflow int32_t");  // as we use int32_t
                                                                               // to store hops
 
-  std::vector<cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>> edge_masks_vector{};
+  std::vector<
+    cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>>
+    edge_masks_vector{};
   graph_view_t<vertex_t, edge_t, false, multi_gpu> modified_graph_view = graph_view;
   edge_masks_vector.reserve(num_edge_types);
-  
+
   if (num_edge_types > 1) {
     for (int i = 0; i < num_edge_types; i++) {
+      cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool>
+        edge_mask(handle, graph_view);
 
-      cugraph::edge_property_t<graph_view_t<vertex_t, edge_t, store_transposed, multi_gpu>, bool> edge_mask(handle, graph_view);
-      
-      cugraph::fill_edge_property(handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
+      cugraph::fill_edge_property(
+        handle, modified_graph_view, edge_mask.mutable_view(), bool{true});
 
       cugraph::transform_e(
         handle,
@@ -118,12 +119,16 @@ neighbor_sample_impl(raft::handle_t const& handle,
         cugraph::edge_src_dummy_property_t{}.view(),
         cugraph::edge_dst_dummy_property_t{}.view(),
         *edge_type_view,
-        [valid_edge_type = i] __device__(auto src, auto dst, thrust::nullopt_t, thrust::nullopt_t, /*thrust::nullopt_t*/auto edge_type) {
+        [valid_edge_type = i] __device__(auto src,
+                                         auto dst,
+                                         thrust::nullopt_t,
+                                         thrust::nullopt_t,
+                                         /*thrust::nullopt_t*/ auto edge_type) {
           return edge_type == valid_edge_type;
         },
         edge_mask.mutable_view(),
         false);
-      
+
       edge_masks_vector.push_back(std::move(edge_mask));
     }
   }
@@ -172,11 +177,13 @@ neighbor_sample_impl(raft::handle_t const& handle,
   std::vector<size_t> level_sizes{};
 
   // Get the number of hop. If homogeneous neighbor sample, num_edge_types = 1
-  auto num_hops = ((fan_out.size() % num_edge_types) == 0) ? (fan_out.size() / num_edge_types) : ((fan_out.size() / num_edge_types) + 1);
+  auto num_hops = ((fan_out.size() % num_edge_types) == 0)
+                    ? (fan_out.size() / num_edge_types)
+                    : ((fan_out.size() / num_edge_types) + 1);
 
-  for (auto hop = 0; hop < num_hops; hop++){
+  for (auto hop = 0; hop < num_hops; hop++) {
     for (auto edge_type_id = 0; edge_type_id < num_edge_types; edge_type_id++) {
-      auto k_level = fan_out[(hop*num_edge_types) + edge_type_id];
+      auto k_level = fan_out[(hop * num_edge_types) + edge_type_id];
       rmm::device_uvector<vertex_t> srcs(0, handle.get_stream());
       rmm::device_uvector<vertex_t> dsts(0, handle.get_stream());
       std::optional<rmm::device_uvector<weight_t>> weights{std::nullopt};
@@ -191,16 +198,16 @@ neighbor_sample_impl(raft::handle_t const& handle,
       if (k_level > 0) {
         std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
           sample_edges(handle,
-                      modified_graph_view,
-                      edge_weight_view,
-                      edge_id_view,
-                      edge_type_view,
-                      edge_bias_view,
-                      rng_state,
-                      starting_vertices,
-                      starting_vertex_labels,
-                      static_cast<size_t>(k_level),
-                      with_replacement);
+                       modified_graph_view,
+                       edge_weight_view,
+                       edge_id_view,
+                       edge_type_view,
+                       edge_bias_view,
+                       rng_state,
+                       starting_vertices,
+                       starting_vertex_labels,
+                       static_cast<size_t>(k_level),
+                       with_replacement);
       } else {
         std::tie(srcs, dsts, weights, edge_ids, edge_types, labels) =
           gather_one_hop_edgelist(handle,
@@ -232,10 +239,10 @@ neighbor_sample_impl(raft::handle_t const& handle,
         starting_vertex_labels,
         raft::device_span<vertex_t const>{level_result_dst_vectors.back().data(),
                                           level_result_dst_vectors.back().size()},
-        frontier_vertex_labels ? std::make_optional(raft::device_span<label_t const>(
-                                    level_result_label_vectors->back().data(),
-                                    level_result_label_vectors->back().size()))
-                                : std::nullopt,
+        frontier_vertex_labels
+          ? std::make_optional(raft::device_span<label_t const>(
+              level_result_label_vectors->back().data(), level_result_label_vectors->back().size()))
+          : std::nullopt,
         std::move(vertex_used_as_source),
         modified_graph_view.local_vertex_partition_view(),
         vertex_partition_range_lasts,
@@ -521,29 +528,30 @@ heterogeneous_uniform_neighbor_sample(
   edge_type_t num_edge_types,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
-{ 
-  using bias_t = weight_t; // dummy
-  
-  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
-    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        edge_id_view,
-        edge_type_view,
-        std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
-        starting_vertices,
-        starting_vertex_labels,
-        label_to_output_comm_rank,
-        fan_out,
-        num_edge_types,
-        sampling_flags.return_hops,
-        sampling_flags.with_replacement,
-        sampling_flags.prior_sources_behavior,
-        sampling_flags.dedupe_sources,
-        do_expensive_check);
-  
+{
+  using bias_t = weight_t;  // dummy
+
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets] =
+    detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+      handle,
+      rng_state,
+      graph_view,
+      edge_weight_view,
+      edge_id_view,
+      edge_type_view,
+      std::optional<edge_property_view_t<edge_t, bias_t const*>>{
+        std::nullopt},  // Optional edge_bias_view
+      starting_vertices,
+      starting_vertex_labels,
+      label_to_output_comm_rank,
+      fan_out,
+      num_edge_types,
+      sampling_flags.return_hops,
+      sampling_flags.with_replacement,
+      sampling_flags.prior_sources_behavior,
+      sampling_flags.dedupe_sources,
+      do_expensive_check);
+
   return std::make_tuple(std::move(majors),
                          std::move(minors),
                          std::move(weights),
@@ -551,7 +559,6 @@ heterogeneous_uniform_neighbor_sample(
                          std::move(edge_types),
                          std::move(hops),
                          std::move(offsets));
-  
 }
 
 template <typename vertex_t,
@@ -583,28 +590,27 @@ heterogeneous_biased_neighbor_sample(
   edge_type_t num_edge_types,
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
-{ 
-  
-  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
-    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        edge_id_view,
-        edge_type_view,
-        std::make_optional(edge_bias_view),
-        starting_vertices,
-        starting_vertex_labels,
-        label_to_output_comm_rank,
-        fan_out,
-        num_edge_types,
-        sampling_flags.return_hops,
-        sampling_flags.with_replacement,
-        sampling_flags.prior_sources_behavior,
-        sampling_flags.dedupe_sources,
-        do_expensive_check);
-  
+{
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets] =
+    detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+      handle,
+      rng_state,
+      graph_view,
+      edge_weight_view,
+      edge_id_view,
+      edge_type_view,
+      std::make_optional(edge_bias_view),
+      starting_vertices,
+      starting_vertex_labels,
+      label_to_output_comm_rank,
+      fan_out,
+      num_edge_types,
+      sampling_flags.return_hops,
+      sampling_flags.with_replacement,
+      sampling_flags.prior_sources_behavior,
+      sampling_flags.dedupe_sources,
+      do_expensive_check);
+
   return std::make_tuple(std::move(majors),
                          std::move(minors),
                          std::move(weights),
@@ -641,28 +647,29 @@ homogeneous_uniform_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
-  using bias_t = weight_t; // dummy
+  using bias_t = weight_t;  // dummy
+
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets] =
+    detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+      handle,
+      rng_state,
+      graph_view,
+      edge_weight_view,
+      edge_id_view,
+      edge_type_view,
+      std::optional<edge_property_view_t<edge_t, bias_t const*>>{
+        std::nullopt},  // Optional edge_bias_view
+      starting_vertices,
+      starting_vertex_labels,
+      label_to_output_comm_rank,
+      fan_out,
+      edge_type_t{1},
+      sampling_flags.return_hops,
+      sampling_flags.with_replacement,
+      sampling_flags.prior_sources_behavior,
+      sampling_flags.dedupe_sources,
+      do_expensive_check);
 
-  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
-    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    rng_state,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    std::optional<edge_property_view_t<edge_t, bias_t const*>>{std::nullopt}, // Optional edge_bias_view
-    starting_vertices,
-    starting_vertex_labels,
-    label_to_output_comm_rank,
-    fan_out,
-    edge_type_t{1},
-    sampling_flags.return_hops,
-    sampling_flags.with_replacement,
-    sampling_flags.prior_sources_behavior,
-    sampling_flags.dedupe_sources,
-    do_expensive_check);
-  
   return std::make_tuple(std::move(majors),
                          std::move(minors),
                          std::move(weights),
@@ -701,27 +708,26 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check)
 {
+  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets] =
+    detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
+      handle,
+      rng_state,
+      graph_view,
+      edge_weight_view,
+      edge_id_view,
+      edge_type_view,
+      std::make_optional(edge_bias_view),
+      starting_vertices,
+      starting_vertex_labels,
+      label_to_output_comm_rank,
+      fan_out,
+      edge_type_t{1},
+      sampling_flags.return_hops,
+      sampling_flags.with_replacement,
+      sampling_flags.prior_sources_behavior,
+      sampling_flags.dedupe_sources,
+      do_expensive_check);
 
-  auto [majors, minors, weights, edge_ids, edge_types, hops, labels, offsets]
-    = detail::neighbor_sample_impl<vertex_t, edge_t, weight_t, edge_type_t, bias_t>(
-    handle,
-    rng_state,
-    graph_view,
-    edge_weight_view,
-    edge_id_view,
-    edge_type_view,
-    std::make_optional(edge_bias_view),
-    starting_vertices,
-    starting_vertex_labels,
-    label_to_output_comm_rank,
-    fan_out,
-    edge_type_t{1},
-    sampling_flags.return_hops,
-    sampling_flags.with_replacement,
-    sampling_flags.prior_sources_behavior,
-    sampling_flags.dedupe_sources,
-    do_expensive_check);
-  
   return std::make_tuple(std::move(majors),
                          std::move(minors),
                          std::move(weights),
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
index b96a73b722d..d848935cc7e 100644
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_mg_v32_e32.cu
@@ -128,12 +128,12 @@ biased_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<float>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -150,12 +150,12 @@ heterogeneous_uniform_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<double>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -218,12 +218,12 @@ heterogeneous_biased_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<float>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -239,12 +239,12 @@ homogeneous_uniform_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<double>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -303,7 +303,4 @@ homogeneous_biased_neighbor_sample(
   sampling_flags_t sampling_flags,
   bool do_expensive_check);
 
-
-
-
 }  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
index 5cb3fd75ce6..72bbb4e27a8 100644
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
+++ b/cpp/src/sampling/neighbor_sampling_sg_v32_e32.cu
@@ -128,12 +128,12 @@ biased_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<float>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -150,12 +150,12 @@ heterogeneous_uniform_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<double>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 heterogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -218,12 +218,12 @@ heterogeneous_biased_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<float>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
@@ -239,12 +239,12 @@ homogeneous_uniform_neighbor_sample(
   bool do_expensive_check);
 
 template std::tuple<rmm::device_uvector<int32_t>,
-                  rmm::device_uvector<int32_t>,
-                  std::optional<rmm::device_uvector<double>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<int32_t>>,
-                  std::optional<rmm::device_uvector<size_t>>>
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
 homogeneous_uniform_neighbor_sample(
   raft::handle_t const& handle,
   raft::random::RngState& rng_state,
diff --git a/cpp/src/utilities/shuffle_vertices.cuh b/cpp/src/utilities/shuffle_vertices.cuh
index 8707068bb90..59b5ecad059 100644
--- a/cpp/src/utilities/shuffle_vertices.cuh
+++ b/cpp/src/utilities/shuffle_vertices.cuh
@@ -45,40 +45,42 @@ rmm::device_uvector<vertex_t> shuffle_vertices_by_gpu_id_impl(
 }
 
 template <typename vertex_t, typename value0_t, typename value1_t, typename func_t>
-std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, std::optional<rmm::device_uvector<value1_t>>>
-shuffle_vertices_and_values_by_gpu_id_impl(raft::handle_t const& handle,
-                                           rmm::device_uvector<vertex_t>&& d_vertices,
-                                           rmm::device_uvector<value0_t>&& d_values_0,
-                                           std::optional<rmm::device_uvector<value1_t>>&& d_values_1,
-                                           func_t func)
+std::tuple<rmm::device_uvector<vertex_t>,
+           rmm::device_uvector<value0_t>,
+           std::optional<rmm::device_uvector<value1_t>>>
+shuffle_vertices_and_values_by_gpu_id_impl(
+  raft::handle_t const& handle,
+  rmm::device_uvector<vertex_t>&& d_vertices,
+  rmm::device_uvector<value0_t>&& d_values_0,
+  std::optional<rmm::device_uvector<value1_t>>&& d_values_1,
+  func_t func)
 {
-
- if (d_values_1) {
-  auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
-    handle.get_comms(),
-    d_vertices.begin(),
-    d_vertices.end(),
-    thrust::make_zip_iterator(d_values_0.begin(), (*d_values_1).begin()),
-    [key_func = func] __device__(auto val) { return key_func(val); },
-    handle.get_stream());
-  
-  return std::make_tuple(
-    std::move(d_shuffled_vertices), std::move(std::get<0>(d_values)), std::make_optional(std::move(std::get<1>(d_values))));
- } else {
-  auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
-    handle.get_comms(),
-    d_vertices.begin(),
-    d_vertices.end(),
-    d_values_0.begin(),
-    [key_func = func] __device__(auto val) { return key_func(val); },
-    handle.get_stream());
+  if (d_values_1) {
+    auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
+      handle.get_comms(),
+      d_vertices.begin(),
+      d_vertices.end(),
+      thrust::make_zip_iterator(d_values_0.begin(), (*d_values_1).begin()),
+      [key_func = func] __device__(auto val) { return key_func(val); },
+      handle.get_stream());
+
+    return std::make_tuple(std::move(d_shuffled_vertices),
+                           std::move(std::get<0>(d_values)),
+                           std::make_optional(std::move(std::get<1>(d_values))));
+  } else {
+    auto [d_shuffled_vertices, d_values, counts] = cugraph::groupby_gpu_id_and_shuffle_kv_pairs(
+      handle.get_comms(),
+      d_vertices.begin(),
+      d_vertices.end(),
+      d_values_0.begin(),
+      [key_func = func] __device__(auto val) { return key_func(val); },
+      handle.get_stream());
 
     auto d_values_1 = std::optional<rmm::device_uvector<int32_t>>{std::nullopt};
 
-  return std::make_tuple(
-    std::move(d_shuffled_vertices), std::move(d_values), std::move(d_values_1));
- }
-  
+    return std::make_tuple(
+      std::move(d_shuffled_vertices), std::move(d_values), std::move(d_values_1));
+  }
 }
 
 }  // namespace
@@ -130,12 +132,13 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
 }
 
 template <typename vertex_t, typename value0_t, typename value1_t>
-std::tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<vertex_t>&& vertices,
-  rmm::device_uvector<value0_t>&& values_0,
-  rmm::device_uvector<value1_t>&& values_1)
+std::
+  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<vertex_t>&& vertices,
+    rmm::device_uvector<value0_t>&& values_0,
+    rmm::device_uvector<value1_t>&& values_1)
 {
   auto const comm_size       = handle.get_comms().get_size();
   auto& major_comm           = handle.get_subcomm(cugraph::partition_manager::major_comm_name());
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
index f148f71f166..c66a0868b6f 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
@@ -40,19 +40,21 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
-template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& vertices,
-  rmm::device_uvector<float>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
+template std::
+  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int32_t>&& vertices,
+    rmm::device_uvector<float>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
 
-template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& vertices,
-  rmm::device_uvector<double>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
+template std::
+  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int32_t>&& vertices,
+    rmm::device_uvector<double>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
 
 }  // namespace detail
 
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
index c8f158bde6a..984f2afc002 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
@@ -62,20 +62,21 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
-template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& vertices,
-  rmm::device_uvector<int32_t>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
-
-template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& vertices,
-  rmm::device_uvector<size_t>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
-
+template std::
+  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int32_t>&& vertices,
+    rmm::device_uvector<int32_t>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
+
+template std::
+  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int32_t>&& vertices,
+    rmm::device_uvector<size_t>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
 
 }  // namespace detail
 
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
index a61c91b8ac3..ff200fac563 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
@@ -40,19 +40,21 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
-template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int64_t>&& vertices,
-  rmm::device_uvector<float>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
+template std::
+  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int64_t>&& vertices,
+    rmm::device_uvector<float>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
 
-template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int64_t>&& vertices,
-  rmm::device_uvector<double>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
+template std::
+  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int64_t>&& vertices,
+    rmm::device_uvector<double>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
 
 }  // namespace detail
 
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
index 1539244307a..fb4c0d469f7 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
@@ -54,19 +54,21 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
-template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int64_t>&& vertices,
-  rmm::device_uvector<int32_t>&& values_0,
-  rmm::device_uvector<int32_t>&& values_1);
-
-template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
-shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int64_t>&& vertices,
-  rmm::device_uvector<size_t>&& values,
-  rmm::device_uvector<int32_t>&& values_1);
+template std::
+  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int64_t>&& vertices,
+    rmm::device_uvector<int32_t>&& values_0,
+    rmm::device_uvector<int32_t>&& values_1);
+
+template std::
+  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
+  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
+    raft::handle_t const& handle,
+    rmm::device_uvector<int64_t>&& vertices,
+    rmm::device_uvector<size_t>&& values,
+    rmm::device_uvector<int32_t>&& values_1);
 
 template rmm::device_uvector<int64_t> shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
   raft::handle_t const& handle,
diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index 04bd3552190..f392c7fb461 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -773,17 +773,17 @@ if(BUILD_CUGRAPH_MG_TESTS)
     # - MG HOMOGENEOUS UNIFORM NBR SAMPLING tests -------------------------------------------------
     ConfigureTestMG(
         MG_HOMOGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/mg_homogeneous_uniform_neighbor_sampling.cpp)
-    
+
     ###############################################################################################
     # - MG HETEROGENEOUS UNIFORM NBR SAMPLING tests -------------------------------------------------
     ConfigureTestMG(
         MG_HETEROGENEOUS_UNIFORM_NEIGHBOR_SAMPLING_TEST sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp)
-    
+
     ###############################################################################################
     # - MG HOMOGENEOUS BIASED NBR SAMPLING tests --------------------------------------------------
     ConfigureTestMG(
         MG_HOMOGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/mg_homogeneous_biased_neighbor_sampling.cpp)
-    
+
     ###############################################################################################
     # - MG HETEROGENEOUS BIASED NBR SAMPLING tests --------------------------------------------------
     ConfigureTestMG(
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index 745a9049c93..c8bfb654598 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -46,8 +46,8 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(
-    std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  void run_current_test(std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase const&,
+                                   input_usecase_t const&> const& param)
   {
     using edge_type_t = int32_t;
 
@@ -101,8 +101,11 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
 
     auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
 
-    batch_number = cugraph::test::sequence(
-      handle, random_sources.size(), heterogeneous_biased_neighbor_sampling_usecase.batch_size, int32_t{0});
+    batch_number =
+      cugraph::test::sequence(handle,
+                              random_sources.size(),
+                              heterogeneous_biased_neighbor_sampling_usecase.batch_size,
+                              int32_t{0});
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
 
@@ -111,9 +114,8 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
                random_sources.size(),
                handle.get_stream());
 
-    std::optional<raft::device_span<int32_t const>>
-      label_to_output_comm_rank_mapping{std::nullopt};
-    
+    std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
+
     // Generate the edge types
 
     std::optional<cugraph::edge_property_t<decltype(graph_view), edge_type_t>> edge_types{
@@ -121,9 +123,7 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
 
     if (heterogeneous_biased_neighbor_sampling_usecase.num_edge_types > 1) {
       edge_types = cugraph::test::generate<decltype(graph_view), edge_type_t>::edge_property(
-          handle,
-          graph_view,
-          heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
+        handle, graph_view, heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
     }
 
 #ifdef NO_CUGRAPH_OPS
@@ -135,25 +135,23 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         *edge_weight_view,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
                      : std::nullopt,
         label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        raft::host_span<int32_t const>(
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_biased_neighbor_sampling_usecase.flag_replacement
-        }
-        ),
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  heterogeneous_biased_neighbor_sampling_usecase.flag_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -169,26 +167,24 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         *edge_weight_view,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
                      : std::nullopt,
         label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        raft::host_span<int32_t const>(
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_biased_neighbor_sampling_usecase.flag_replacement
-        }
-        );
-    
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  heterogeneous_biased_neighbor_sampling_usecase.flag_replacement});
+
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.stop();
@@ -224,22 +220,19 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
         raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
         true);
-      
-      
-      
 
       ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
         handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
 
       if (random_sources.size() < 100) {
         // This validation is too expensive for large number of vertices
-        ASSERT_TRUE(
-          cugraph::test::validate_sampling_depth(handle,
-                                                 std::move(src_out),
-                                                 std::move(dst_out),
-                                                 std::move(wgt_out),
-                                                 std::move(random_sources),
-                                                 heterogeneous_biased_neighbor_sampling_usecase.fanout.size()));
+        ASSERT_TRUE(cugraph::test::validate_sampling_depth(
+          handle,
+          std::move(src_out),
+          std::move(dst_out),
+          std::move(wgt_out),
+          std::move(random_sources),
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()));
       }
     }
 #endif
@@ -288,7 +281,6 @@ TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-
 INSTANTIATE_TEST_SUITE_P(
   file_test,
   Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
@@ -329,11 +321,16 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
+    ::testing::Values(
+      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
+      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
-//#endif
+// #endif
 
 CUGRAPH_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index d94d7ad7e5a..15be79a623e 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -46,8 +46,8 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(
-    std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  void run_current_test(std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase const&,
+                                   input_usecase_t const&> const& param)
   {
     using edge_type_t = int32_t;
 
@@ -99,8 +99,11 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
     auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
 
-    batch_number = cugraph::test::sequence(
-      handle, random_sources.size(), heterogeneous_uniform_neighbor_sampling_usecase.batch_size, int32_t{0});
+    batch_number =
+      cugraph::test::sequence(handle,
+                              random_sources.size(),
+                              heterogeneous_uniform_neighbor_sampling_usecase.batch_size,
+                              int32_t{0});
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
 
@@ -109,9 +112,8 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
                random_sources.size(),
                handle.get_stream());
 
-    std::optional<raft::device_span<int32_t const>>
-      label_to_output_comm_rank_mapping{std::nullopt};
-    
+    std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
+
     // Generate the edge types
 
     std::optional<cugraph::edge_property_t<decltype(graph_view), edge_type_t>> edge_types{
@@ -119,9 +121,7 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
     if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
       edge_types = cugraph::test::generate<decltype(graph_view), edge_type_t>::edge_property(
-          handle,
-          graph_view,
-          heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
+        handle, graph_view, heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
     }
 
 #ifdef NO_CUGRAPH_OPS
@@ -133,24 +133,23 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
                      : std::nullopt,
         label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        raft::host_span<int32_t const>(
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
         cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement
-        }
-        ),
+          cugraph::prior_sources_behavior_t{0},
+          true,   // return_hops
+          false,  // dedupe_sources
+          heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -166,25 +165,24 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
         edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
         batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
                                                                            batch_number->size()})
                      : std::nullopt,
         label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        raft::host_span<int32_t const>(
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
         cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement
-        }
-        );
-    
+          cugraph::prior_sources_behavior_t{0},
+          true,   // return_hops
+          false,  // dedupe_sources
+          heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement});
+
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.stop();
@@ -220,22 +218,19 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
         raft::device_span<size_t const>(d_subgraph_offsets.data(), 2),
         raft::device_span<vertex_t const>(vertices.data(), vertices.size()),
         true);
-      
-      
-      
 
       ASSERT_TRUE(cugraph::test::validate_extracted_graph_is_subgraph(
         handle, src_compare, dst_compare, wgt_compare, src_out, dst_out, wgt_out));
 
       if (random_sources.size() < 100) {
         // This validation is too expensive for large number of vertices
-        ASSERT_TRUE(
-          cugraph::test::validate_sampling_depth(handle,
-                                                 std::move(src_out),
-                                                 std::move(dst_out),
-                                                 std::move(wgt_out),
-                                                 std::move(random_sources),
-                                                 heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
+        ASSERT_TRUE(cugraph::test::validate_sampling_depth(
+          handle,
+          std::move(src_out),
+          std::move(dst_out),
+          std::move(wgt_out),
+          std::move(random_sources),
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
       }
     }
 #endif
@@ -248,7 +243,6 @@ using Tests_Heterogeneous_Uniform_Neighbor_Sampling_File =
 using Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat =
   Tests_Heterogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
 
-
 TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
   run_current_test<int32_t, int32_t, float>(
@@ -285,7 +279,6 @@ TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-
 INSTANTIATE_TEST_SUITE_P(
   file_test,
   Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
@@ -326,11 +319,16 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
+    ::testing::Values(
+      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
+      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
+      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
+        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
-//#endif
+// #endif
 
 CUGRAPH_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
index a3f6b612337..d9f3e90f127 100644
--- a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -46,8 +46,8 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(
-    std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  void run_current_test(std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase const&,
+                                   input_usecase_t const&> const& param)
   {
     auto [homogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
 
@@ -104,8 +104,10 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
 
     auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
 
-    batch_number = cugraph::test::sequence(
-      handle, random_sources.size(), homogeneous_biased_neighbor_sampling_usecase.batch_size, int32_t{0});
+    batch_number = cugraph::test::sequence(handle,
+                                           random_sources.size(),
+                                           homogeneous_biased_neighbor_sampling_usecase.batch_size,
+                                           int32_t{0});
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
 
@@ -114,8 +116,7 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
                random_sources.size(),
                handle.get_stream());
 
-    std::optional<raft::device_span<int32_t const>>
-      label_to_output_comm_rank_mapping{std::nullopt};
+    std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
 
 #ifdef NO_CUGRAPH_OPS
     EXPECT_THROW(
@@ -134,13 +135,10 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_biased_neighbor_sampling_usecase.flag_replacement
-        }
-        ),
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_biased_neighbor_sampling_usecase.flag_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -164,14 +162,10 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_biased_neighbor_sampling_usecase.flag_replacement
-        }
-        );
-    
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_biased_neighbor_sampling_usecase.flag_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -214,13 +208,13 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
 
       if (random_sources.size() < 100) {
         // This validation is too expensive for large number of vertices
-        ASSERT_TRUE(
-          cugraph::test::validate_sampling_depth(handle,
-                                                 std::move(src_out),
-                                                 std::move(dst_out),
-                                                 std::move(wgt_out),
-                                                 std::move(random_sources),
-                                                 homogeneous_biased_neighbor_sampling_usecase.fanout.size()));
+        ASSERT_TRUE(cugraph::test::validate_sampling_depth(
+          handle,
+          std::move(src_out),
+          std::move(dst_out),
+          std::move(wgt_out),
+          std::move(random_sources),
+          homogeneous_biased_neighbor_sampling_usecase.fanout.size()));
       }
     }
 #endif
@@ -233,7 +227,6 @@ using Tests_Homogeneous_Biased_Neighbor_Sampling_File =
 using Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat =
   Tests_Homogeneous_Biased_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
 
-
 TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
   run_current_test<int32_t, int32_t, float>(
@@ -310,10 +303,11 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
+    ::testing::Values(
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
 
 CUGRAPH_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
index 5b18ecf6a50..8c50ab56917 100644
--- a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -46,8 +46,8 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(
-    std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase const&, input_usecase_t const&> const& param)
+  void run_current_test(std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase const&,
+                                   input_usecase_t const&> const& param)
   {
     auto [homogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
 
@@ -104,8 +104,10 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
 
     auto batch_number = std::make_optional<rmm::device_uvector<int32_t>>(0, handle.get_stream());
 
-    batch_number = cugraph::test::sequence(
-      handle, random_sources.size(), homogeneous_uniform_neighbor_sampling_usecase.batch_size, int32_t{0});
+    batch_number = cugraph::test::sequence(handle,
+                                           random_sources.size(),
+                                           homogeneous_uniform_neighbor_sampling_usecase.batch_size,
+                                           int32_t{0});
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle.get_stream());
 
@@ -114,8 +116,7 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
                random_sources.size(),
                handle.get_stream());
 
-    std::optional<raft::device_span<int32_t const>>
-      label_to_output_comm_rank_mapping{std::nullopt};
+    std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
 
 #ifdef NO_CUGRAPH_OPS
     EXPECT_THROW(
@@ -133,13 +134,10 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_uniform_neighbor_sampling_usecase.flag_replacement
-        }
-        ),
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_uniform_neighbor_sampling_usecase.flag_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -162,14 +160,10 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
         label_to_output_comm_rank_mapping,
         raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_uniform_neighbor_sampling_usecase.flag_replacement
-        }
-        );
-    
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_uniform_neighbor_sampling_usecase.flag_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -212,13 +206,13 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
 
       if (random_sources.size() < 100) {
         // This validation is too expensive for large number of vertices
-        ASSERT_TRUE(
-          cugraph::test::validate_sampling_depth(handle,
-                                                 std::move(src_out),
-                                                 std::move(dst_out),
-                                                 std::move(wgt_out),
-                                                 std::move(random_sources),
-                                                 homogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
+        ASSERT_TRUE(cugraph::test::validate_sampling_depth(
+          handle,
+          std::move(src_out),
+          std::move(dst_out),
+          std::move(wgt_out),
+          std::move(random_sources),
+          homogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
       }
     }
 #endif
@@ -231,7 +225,6 @@ using Tests_Homogeneous_Uniform_Neighbor_Sampling_File =
 using Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat =
   Tests_Homogeneous_Uniform_Neighbor_Sampling<cugraph::test::Rmat_Usecase>;
 
-
 TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
 {
   run_current_test<int32_t, int32_t, float>(
@@ -308,10 +301,11 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
+    ::testing::Values(
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, false, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, false, true, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, false, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 1024, true, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
 
 CUGRAPH_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index d9cdee3479c..c2e34355062 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -50,7 +50,8 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  void run_current_test(
+    std::tuple<Heterogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
   {
     using edge_type_t = int32_t;
 
@@ -86,7 +87,6 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
     auto mg_edge_weight_view =
       mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
 
-
     //
     // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
     // batches
@@ -101,11 +101,11 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
       mg_graph_view,
       std::optional<raft::device_span<vertex_t const>>{std::nullopt},
       rng_state,
-      //20,
-      
+      // 20,
+
       std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
                std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
-      
+
       false,
       false);
 
@@ -115,24 +115,22 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
 
     auto seed_sizes = cugraph::host_scalar_allgather(
       handle_->get_comms(), random_sources.size(), handle_->get_stream());
-    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
-    size_t num_batches = (num_seeds + heterogeneous_biased_neighbor_sampling_usecase.batch_size - 1) /
-                         heterogeneous_biased_neighbor_sampling_usecase.batch_size;
+    size_t num_seeds = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches =
+      (num_seeds + heterogeneous_biased_neighbor_sampling_usecase.batch_size - 1) /
+      heterogeneous_biased_neighbor_sampling_usecase.batch_size;
 
     std::vector<size_t> seed_offsets(seed_sizes.size());
     std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
 
     auto batch_number = cugraph::test::modulo_sequence<int32_t>(
       *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
-    
+
     // Get unique batch_number -> label_list
     rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
 
-    raft::copy(label_list.data(),
-               batch_number.data(),
-               batch_number.size(),
-               handle_->get_stream());
-    
+    raft::copy(label_list.data(), batch_number.data(), batch_number.size(), handle_->get_stream());
+
     label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
     label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
 
@@ -142,8 +140,7 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
       *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
 
     // perform allgatherv
-    comm_ranks =
-      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+    comm_ranks = cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
 
@@ -151,7 +148,7 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
                random_sources.data(),
                random_sources.size(),
                handle_->get_stream());
-    
+
     // Generate the edge types
 
     std::optional<cugraph::edge_property_t<decltype(mg_graph_view), edge_type_t>> edge_types{
@@ -159,9 +156,7 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
 
     if (heterogeneous_biased_neighbor_sampling_usecase.num_edge_types > 1) {
       edge_types = cugraph::test::generate<decltype(mg_graph_view), edge_type_t>::edge_property(
-          *handle_,
-          mg_graph_view,
-          heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
+        *handle_, mg_graph_view, heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
     }
 
 #ifdef NO_CUGRAPH_OPS
@@ -173,25 +168,22 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
-        
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
+
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_biased_neighbor_sampling_usecase.with_replacement
-        }
-        ),
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  heterogeneous_biased_neighbor_sampling_usecase.with_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -200,7 +192,7 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
       hr_timer.start("MG biased_neighbor_sample");
     }
     RAFT_CUDA_TRY(cudaDeviceSynchronize());
-  
+
     auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
       cugraph::heterogeneous_biased_neighbor_sample(
         *handle_,
@@ -209,25 +201,22 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         *mg_edge_weight_view,
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_biased_neighbor_sampling_usecase.with_replacement
-        }
-        );
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  heterogeneous_biased_neighbor_sampling_usecase.with_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -309,12 +298,13 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
         if (random_sources.size() < 100) {
           // This validation is too expensive for large number of vertices
           if (mg_aggregate_src.size() > 0) {
-            cugraph::test::validate_sampling_depth(*handle_,
-                                                   std::move(mg_aggregate_src),
-                                                   std::move(mg_aggregate_dst),
-                                                   std::move(mg_aggregate_wgt),
-                                                   std::move(mg_start_src),
-                                                   heterogeneous_biased_neighbor_sampling_usecase.fanout.size());
+            cugraph::test::validate_sampling_depth(
+              *handle_,
+              std::move(mg_aggregate_src),
+              std::move(mg_aggregate_dst),
+              std::move(mg_aggregate_wgt),
+              std::move(mg_start_src),
+              heterogeneous_biased_neighbor_sampling_usecase.fanout.size());
           }
         }
       }
@@ -327,8 +317,8 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
 };
 
 template <typename input_usecase_t>
-std::unique_ptr<raft::handle_t> Tests_MGHeterogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ =
-  nullptr;
+std::unique_ptr<raft::handle_t>
+  Tests_MGHeterogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ = nullptr;
 
 using Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File =
   Tests_MGHeterogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index 87dab4801fb..6282f57f1be 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -50,7 +50,8 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  void run_current_test(
+    std::tuple<Heterogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
   {
     using edge_type_t = int32_t;
 
@@ -86,7 +87,6 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
     auto mg_edge_weight_view =
       mg_edge_weights ? std::make_optional((*mg_edge_weights).view()) : std::nullopt;
 
-
     //
     // Test is designed like GNN sampling.  We'll select 5% of vertices to be included in sampling
     // batches
@@ -101,11 +101,11 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
       mg_graph_view,
       std::optional<raft::device_span<vertex_t const>>{std::nullopt},
       rng_state,
-      //20,
-      
+      // 20,
+
       std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
                std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
-      
+
       false,
       false);
 
@@ -115,24 +115,22 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
 
     auto seed_sizes = cugraph::host_scalar_allgather(
       handle_->get_comms(), random_sources.size(), handle_->get_stream());
-    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
-    size_t num_batches = (num_seeds + heterogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
-                         heterogeneous_uniform_neighbor_sampling_usecase.batch_size;
+    size_t num_seeds = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches =
+      (num_seeds + heterogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
+      heterogeneous_uniform_neighbor_sampling_usecase.batch_size;
 
     std::vector<size_t> seed_offsets(seed_sizes.size());
     std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
 
     auto batch_number = cugraph::test::modulo_sequence<int32_t>(
       *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
-    
+
     // Get unique batch_number -> label_list
     rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
 
-    raft::copy(label_list.data(),
-               batch_number.data(),
-               batch_number.size(),
-               handle_->get_stream());
-    
+    raft::copy(label_list.data(), batch_number.data(), batch_number.size(), handle_->get_stream());
+
     label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
     label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
 
@@ -142,8 +140,7 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
       *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
 
     // perform allgatherv
-    comm_ranks =
-      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+    comm_ranks = cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
 
@@ -151,7 +148,7 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
                random_sources.data(),
                random_sources.size(),
                handle_->get_stream());
-    
+
     // Generate the edge types
 
     std::optional<cugraph::edge_property_t<decltype(mg_graph_view), edge_type_t>> edge_types{
@@ -159,9 +156,7 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
 
     if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
       edge_types = cugraph::test::generate<decltype(mg_graph_view), edge_type_t>::edge_property(
-          *handle_,
-          mg_graph_view,
-          heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
+        *handle_, mg_graph_view, heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
     }
 
 #ifdef NO_CUGRAPH_OPS
@@ -173,25 +168,23 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
-        
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
+
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
         cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_uniform_neighbor_sampling_usecase.with_replacement
-        }
-        ),
+          cugraph::prior_sources_behavior_t{0},
+          true,   // return_hops
+          false,  // dedupe_sources
+          heterogeneous_uniform_neighbor_sampling_usecase.with_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -200,7 +193,7 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
       hr_timer.start("MG uniform_neighbor_sample");
     }
     RAFT_CUDA_TRY(cudaDeviceSynchronize());
-  
+
     auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
       cugraph::heterogeneous_uniform_neighbor_sample(
         *handle_,
@@ -209,24 +202,22 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         edge_types
-            ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                        .view()}
-            : std::nullopt,
+          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
+                                                                                       .view()}
+          : std::nullopt,
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        raft::host_span<int32_t const>(
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
+          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
         heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
         cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            heterogeneous_uniform_neighbor_sampling_usecase.with_replacement
-        }
-        );
+          cugraph::prior_sources_behavior_t{0},
+          true,   // return_hops
+          false,  // dedupe_sources
+          heterogeneous_uniform_neighbor_sampling_usecase.with_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -308,12 +299,13 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
         if (random_sources.size() < 100) {
           // This validation is too expensive for large number of vertices
           if (mg_aggregate_src.size() > 0) {
-            cugraph::test::validate_sampling_depth(*handle_,
-                                                   std::move(mg_aggregate_src),
-                                                   std::move(mg_aggregate_dst),
-                                                   std::move(mg_aggregate_wgt),
-                                                   std::move(mg_start_src),
-                                                   heterogeneous_uniform_neighbor_sampling_usecase.fanout.size());
+            cugraph::test::validate_sampling_depth(
+              *handle_,
+              std::move(mg_aggregate_src),
+              std::move(mg_aggregate_dst),
+              std::move(mg_aggregate_wgt),
+              std::move(mg_start_src),
+              heterogeneous_uniform_neighbor_sampling_usecase.fanout.size());
           }
         }
       }
@@ -326,8 +318,8 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
 };
 
 template <typename input_usecase_t>
-std::unique_ptr<raft::handle_t> Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ =
-  nullptr;
+std::unique_ptr<raft::handle_t>
+  Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ = nullptr;
 
 using Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File =
   Tests_MGHeterogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
@@ -404,7 +396,8 @@ INSTANTIATE_TEST_SUITE_P(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
+                        {4, 10, 7, 8}, 128, 2, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
index 1ac6c4244ce..4ad5f0a2053 100644
--- a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
@@ -50,7 +50,8 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  void run_current_test(
+    std::tuple<Homogeneous_Biased_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
   {
     auto [homogeneous_biased_neighbor_sampling_usecase, input_usecase] = param;
 
@@ -105,11 +106,11 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
       mg_graph_view,
       std::optional<raft::device_span<vertex_t const>>{std::nullopt},
       rng_state,
-      //20,
-      
+      // 20,
+
       std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
                std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
-      
+
       false,
       false);
 
@@ -128,15 +129,12 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
 
     auto batch_number = cugraph::test::modulo_sequence<int32_t>(
       *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
-    
+
     // Get unique batch_number -> label_list
     rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
 
-    raft::copy(label_list.data(),
-               batch_number.data(),
-               batch_number.size(),
-               handle_->get_stream());
-    
+    raft::copy(label_list.data(), batch_number.data(), batch_number.size(), handle_->get_stream());
+
     label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
     label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
 
@@ -146,8 +144,7 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
       *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
 
     // perform allgatherv
-    comm_ranks =
-      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+    comm_ranks = cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
 
@@ -165,22 +162,18 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
         mg_edge_weight_view,
         std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
         std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        mg_edge_weight_view
-        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
+        mg_edge_weight_view raft::device_span<vertex_t const>{random_sources.data(),
+                                                              random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
         raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_biased_neighbor_sampling_usecase.with_replacement
-        }
-        ),
+
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_biased_neighbor_sampling_usecase.with_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -201,18 +194,14 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
         raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_biased_neighbor_sampling_usecase.with_replacement
-        }
-        );
+
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_biased_neighbor_sampling_usecase.with_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -294,12 +283,13 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
         if (random_sources.size() < 100) {
           // This validation is too expensive for large number of vertices
           if (mg_aggregate_src.size() > 0) {
-            cugraph::test::validate_sampling_depth(*handle_,
-                                                   std::move(mg_aggregate_src),
-                                                   std::move(mg_aggregate_dst),
-                                                   std::move(mg_aggregate_wgt),
-                                                   std::move(mg_start_src),
-                                                   homogeneous_biased_neighbor_sampling_usecase.fanout.size());
+            cugraph::test::validate_sampling_depth(
+              *handle_,
+              std::move(mg_aggregate_src),
+              std::move(mg_aggregate_dst),
+              std::move(mg_aggregate_wgt),
+              std::move(mg_start_src),
+              homogeneous_biased_neighbor_sampling_usecase.fanout.size());
           }
         }
       }
@@ -312,8 +302,8 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
 };
 
 template <typename input_usecase_t>
-std::unique_ptr<raft::handle_t> Tests_MGHomogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ =
-  nullptr;
+std::unique_ptr<raft::handle_t>
+  Tests_MGHomogeneous_Biased_Neighbor_Sampling<input_usecase_t>::handle_ = nullptr;
 
 using Tests_MGHomogeneous_Biased_Neighbor_Sampling_File =
   Tests_MGHomogeneous_Biased_Neighbor_Sampling<cugraph::test::File_Usecase>;
@@ -387,10 +377,11 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
-                      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
+    ::testing::Values(
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
+      Homogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
index 7132231cf15..68ca7d4dd8e 100644
--- a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -50,7 +50,8 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
   virtual void TearDown() {}
 
   template <typename vertex_t, typename edge_t, typename weight_t>
-  void run_current_test(std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
+  void run_current_test(
+    std::tuple<Homogeneous_Uniform_Neighbor_Sampling_Usecase, input_usecase_t> const& param)
   {
     auto [homogeneous_uniform_neighbor_sampling_usecase, input_usecase] = param;
 
@@ -105,11 +106,11 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
       mg_graph_view,
       std::optional<raft::device_span<vertex_t const>>{std::nullopt},
       rng_state,
-      //20,
-      
+      // 20,
+
       std::max(static_cast<size_t>(mg_graph_view.number_of_vertices() * select_probability),
                std::min(static_cast<size_t>(mg_graph_view.number_of_vertices()), size_t{1})),
-      
+
       false,
       false);
 
@@ -119,24 +120,22 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
 
     auto seed_sizes = cugraph::host_scalar_allgather(
       handle_->get_comms(), random_sources.size(), handle_->get_stream());
-    size_t num_seeds   = std::reduce(seed_sizes.begin(), seed_sizes.end());
-    size_t num_batches = (num_seeds + homogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
-                         homogeneous_uniform_neighbor_sampling_usecase.batch_size;
+    size_t num_seeds = std::reduce(seed_sizes.begin(), seed_sizes.end());
+    size_t num_batches =
+      (num_seeds + homogeneous_uniform_neighbor_sampling_usecase.batch_size - 1) /
+      homogeneous_uniform_neighbor_sampling_usecase.batch_size;
 
     std::vector<size_t> seed_offsets(seed_sizes.size());
     std::exclusive_scan(seed_sizes.begin(), seed_sizes.end(), seed_offsets.begin(), size_t{0});
 
     auto batch_number = cugraph::test::modulo_sequence<int32_t>(
       *handle_, random_sources.size(), num_batches, seed_offsets[handle_->get_comms().get_rank()]);
-    
+
     // Get unique batch_number -> label_list
     rmm::device_uvector<int32_t> label_list(batch_number.size(), handle_->get_stream());
 
-    raft::copy(label_list.data(),
-               batch_number.data(),
-               batch_number.size(),
-               handle_->get_stream());
-    
+    raft::copy(label_list.data(), batch_number.data(), batch_number.size(), handle_->get_stream());
+
     label_list = cugraph::test::sort<int32_t>(*handle_, std::move(label_list));
     label_list = cugraph::test::unique<int32_t>(*handle_, std::move(label_list));
 
@@ -146,8 +145,7 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
       *handle_, num_unique_labels, int32_t{handle_->get_comms().get_rank()});
 
     // perform allgatherv
-    comm_ranks =
-      cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
+    comm_ranks = cugraph::test::device_allgatherv(*handle_, comm_ranks.data(), comm_ranks.size());
 
     rmm::device_uvector<vertex_t> random_sources_copy(random_sources.size(), handle_->get_stream());
 
@@ -168,18 +166,14 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
         raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_uniform_neighbor_sampling_usecase.with_replacement
-        }
-        ),
+
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_uniform_neighbor_sampling_usecase.with_replacement}),
       std::exception);
 #else
     if (cugraph::test::g_perf) {
@@ -188,7 +182,7 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
       hr_timer.start("MG uniform_neighbor_sample");
     }
     RAFT_CUDA_TRY(cudaDeviceSynchronize());
-  
+
     auto&& [src_out, dst_out, wgt_out, edge_id, edge_type, hop, offsets] =
       cugraph::homogeneous_uniform_neighbor_sample(
         *handle_,
@@ -200,18 +194,14 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
         raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
         std::make_optional(
           raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(
-          raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
+        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
         raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
                                        homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        
-        cugraph::sampling_flags_t{
-            cugraph::prior_sources_behavior_t{0},
-            true, // return_hops
-            false, // dedupe_sources
-            homogeneous_uniform_neighbor_sampling_usecase.with_replacement
-        }
-        );
+
+        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
+                                  true,   // return_hops
+                                  false,  // dedupe_sources
+                                  homogeneous_uniform_neighbor_sampling_usecase.with_replacement});
 
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
@@ -293,12 +283,13 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
         if (random_sources.size() < 100) {
           // This validation is too expensive for large number of vertices
           if (mg_aggregate_src.size() > 0) {
-            cugraph::test::validate_sampling_depth(*handle_,
-                                                   std::move(mg_aggregate_src),
-                                                   std::move(mg_aggregate_dst),
-                                                   std::move(mg_aggregate_wgt),
-                                                   std::move(mg_start_src),
-                                                   homogeneous_uniform_neighbor_sampling_usecase.fanout.size());
+            cugraph::test::validate_sampling_depth(
+              *handle_,
+              std::move(mg_aggregate_src),
+              std::move(mg_aggregate_dst),
+              std::move(mg_aggregate_wgt),
+              std::move(mg_start_src),
+              homogeneous_uniform_neighbor_sampling_usecase.fanout.size());
           }
         }
       }
@@ -311,8 +302,8 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
 };
 
 template <typename input_usecase_t>
-std::unique_ptr<raft::handle_t> Tests_MGHomogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ =
-  nullptr;
+std::unique_ptr<raft::handle_t>
+  Tests_MGHomogeneous_Uniform_Neighbor_Sampling<input_usecase_t>::handle_ = nullptr;
 
 using Tests_MGHomogeneous_Uniform_Neighbor_Sampling_File =
   Tests_MGHomogeneous_Uniform_Neighbor_Sampling<cugraph::test::File_Usecase>;
@@ -386,10 +377,11 @@ INSTANTIATE_TEST_SUITE_P(
                           factor (to avoid running same benchmarks more than once) */
   Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
-    ::testing::Values(Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
-                      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
+    ::testing::Values(
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, false, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, false, true, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, false, false},
+      Homogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10}, 128, true, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/utilities/thrust_wrapper.cu b/cpp/tests/utilities/thrust_wrapper.cu
index 2d759a22b1f..539de42845f 100644
--- a/cpp/tests/utilities/thrust_wrapper.cu
+++ b/cpp/tests/utilities/thrust_wrapper.cu
@@ -41,7 +41,7 @@ namespace test {
 
 template <typename value_t>
 cugraph::dataframe_buffer_type_t<value_t> sort(
-  
+
   raft::handle_t const& handle, cugraph::dataframe_buffer_type_t<value_t> const& values)
 {
   auto sorted_values = cugraph::allocate_dataframe_buffer<value_t>(
@@ -410,14 +410,12 @@ cugraph::dataframe_buffer_type_t<value_t> scalar_fill(raft::handle_t const& hand
                                                       value_t value)
 {
   auto values = cugraph::allocate_dataframe_buffer<value_t>(length, handle.get_stream());
-  
-  thrust::tabulate(handle.get_thrust_policy(),
-                    values.begin(),
-                    values.end(),
-                    [value] __device__(size_t i) {
-                      return value;
-                    });
-  
+
+  thrust::tabulate(
+    handle.get_thrust_policy(), values.begin(), values.end(), [value] __device__(size_t i) {
+      return value;
+    });
+
   return values;
 }
 
diff --git a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index 45805b88581..36a86635254 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -620,7 +620,7 @@ def uniform_neighbor_sample(
         # FIXME: Add expensive check to ensure all dict values are lists
         # Convert to a tuple of sequence (edge type size and fanout values)
 
-        fanout_vals = {2:[1], 3:[4, 5], 1:[6, 7], 0:[8, 9, 10]}
+        fanout_vals = {2: [1], 3: [4, 5], 1: [6, 7], 0: [8, 9, 10]}
         fanout_vals = collections.OrderedDict(sorted(fanout_vals.items()))
 
         edge_type_offsets = []
diff --git a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
index b953a5d14f4..0266bb6ddb6 100644
--- a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
@@ -14,8 +14,9 @@
 from __future__ import annotations
 
 from pylibcugraph import ResourceHandle
-from pylibcugraph import heterogeneous_biased_neighbor_sample as \
-    pylibcugraph_heterogeneous_biased_neighbor_sample
+from pylibcugraph import (
+    heterogeneous_biased_neighbor_sample as pylibcugraph_heterogeneous_biased_neighbor_sample,
+)
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
 
@@ -87,14 +88,14 @@ def heterogeneous_biased_neighbor_sample(
 
     start_vertex_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     fanout_vals : list (int32)
         List of branching out (fan-out) degrees per starting vertex for each
         hop level.
-    
+
     num_edge_types: int32
         Number of edge types where a value of 1 translates to homogeneous neighbor
         sample whereas a value greater than 1 translates to heterogeneous neighbor
@@ -107,7 +108,7 @@ def heterogeneous_biased_neighbor_sample(
         Deprecated.
         Flag to specify whether to return edge properties (weight, edge id,
         edge type, batch id, hop id) with the sampled edges.
-    
+
     prior_sources_behavior: str, optional (default=None)
         Options are "carryover", and "exclude".
         Default will leave the source list as-is.
@@ -115,37 +116,37 @@ def heterogeneous_biased_neighbor_sample(
         current hop.
         Exclude will exclude sources from previous hops from reappearing
         as sources in future hops.
-    
+
     deduplicate_sources: bool, optional (default=False)
         Whether to first deduplicate the list of possible sources
         from the previous destinations before performing next
         hop.
-    
+
     return_hops: bool, optional (default=True)
         Whether to return the sampling results with hop ids
         corresponding to the hop where the edge appeared.
         Defaults to True.
-    
+
     renumber: bool, optional (default=False)
         Whether to renumber on a per-batch basis.  If True,
         will return the renumber map and renumber map offsets
         as an additional dataframe.
-    
+
     retain_seeds: bool, optional (default=False)
         If True, will retain the original seeds (original source vertices)
         in the output even if they do not have outgoing neighbors.
-    
+
     compression: str, optional (default=COO)
         Sets the compression type for the output minibatches.
         Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-    
+
     compress_per_hop: bool, optional (default=False)
         Whether to compress globally (default), or to produce a separate
         compressed edgelist per hop.
 
     random_state: int, optional
         Random seed to use when making sampling calls.
-    
+
     return_offsets: bool, optional (default=False)
         Whether to return the sampling results with batch ids
         included as one dataframe, or to instead return two
@@ -216,8 +217,8 @@ def heterogeneous_biased_neighbor_sample(
                         Contains the batch offsets for the renumber maps
     """
 
-    use_legacy_names = False # Deprecated parameter
-    include_hop_column=not return_offsets # Deprecated parameter
+    use_legacy_names = False  # Deprecated parameter
+    include_hop_column = not return_offsets  # Deprecated parameter
 
     major_col_name = "majors"
     minor_col_name = "minors"
@@ -243,7 +244,6 @@ def heterogeneous_biased_neighbor_sample(
             " of the libcugraph C++ API"
         )
 
-
     if with_edge_properties:
         warning_msg = (
             "The with_edge_properties flag is deprecated"
@@ -260,8 +260,7 @@ def heterogeneous_biased_neighbor_sample(
             start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
         )
 
-
-    """ 
+    """
     # No batch_ids, the rank owning the vertices will wom the final
     # result.
     if with_edge_properties and not with_batch_ids:
@@ -292,11 +291,9 @@ def heterogeneous_biased_neighbor_sample(
 
     start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
 
-
     if G.renumbered:
         start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
 
-
     sampling_result_array_dict = pylibcugraph_heterogeneous_biased_neighbor_sample(
         resource_handle=ResourceHandle(),
         input_graph=G._plc_graph,
@@ -326,7 +323,7 @@ def heterogeneous_biased_neighbor_sample(
         return_offsets=return_offsets,
         renumber=renumber,
         use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column, # Deprecated flag
+        include_hop_column=include_hop_column,  # Deprecated flag
     )
 
     if G.renumbered and not renumber:
diff --git a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
index 7eadebbc912..3c81877abc6 100644
--- a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
@@ -14,8 +14,9 @@
 from __future__ import annotations
 
 from pylibcugraph import ResourceHandle
-from pylibcugraph import heterogeneous_uniform_neighbor_sample as \
-    pylibcugraph_heterogeneous_uniform_neighbor_sample
+from pylibcugraph import (
+    heterogeneous_uniform_neighbor_sample as pylibcugraph_heterogeneous_uniform_neighbor_sample,
+)
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
 
@@ -87,14 +88,14 @@ def heterogeneous_uniform_neighbor_sample(
 
     start_vertex_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     fanout_vals : list (int32)
         List of branching out (fan-out) degrees per starting vertex for each
         hop level.
-    
+
     num_edge_types: int32
         Number of edge types where a value of 1 translates to homogeneous neighbor
         sample whereas a value greater than 1 translates to heterogeneous neighbor
@@ -107,7 +108,7 @@ def heterogeneous_uniform_neighbor_sample(
         Deprecated.
         Flag to specify whether to return edge properties (weight, edge id,
         edge type, batch id, hop id) with the sampled edges.
-    
+
     prior_sources_behavior: str, optional (default=None)
         Options are "carryover", and "exclude".
         Default will leave the source list as-is.
@@ -115,37 +116,37 @@ def heterogeneous_uniform_neighbor_sample(
         current hop.
         Exclude will exclude sources from previous hops from reappearing
         as sources in future hops.
-    
+
     deduplicate_sources: bool, optional (default=False)
         Whether to first deduplicate the list of possible sources
         from the previous destinations before performing next
         hop.
-    
+
     return_hops: bool, optional (default=True)
         Whether to return the sampling results with hop ids
         corresponding to the hop where the edge appeared.
         Defaults to True.
-    
+
     renumber: bool, optional (default=False)
         Whether to renumber on a per-batch basis.  If True,
         will return the renumber map and renumber map offsets
         as an additional dataframe.
-    
+
     retain_seeds: bool, optional (default=False)
         If True, will retain the original seeds (original source vertices)
         in the output even if they do not have outgoing neighbors.
-    
+
     compression: str, optional (default=COO)
         Sets the compression type for the output minibatches.
         Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-    
+
     compress_per_hop: bool, optional (default=False)
         Whether to compress globally (default), or to produce a separate
         compressed edgelist per hop.
 
     random_state: int, optional
         Random seed to use when making sampling calls.
-    
+
     return_offsets: bool, optional (default=False)
         Whether to return the sampling results with batch ids
         included as one dataframe, or to instead return two
@@ -216,8 +217,8 @@ def heterogeneous_uniform_neighbor_sample(
                         Contains the batch offsets for the renumber maps
     """
 
-    use_legacy_names = False # Deprecated parameter
-    include_hop_column=not return_offsets # Deprecated parameter
+    use_legacy_names = False  # Deprecated parameter
+    include_hop_column = not return_offsets  # Deprecated parameter
 
     major_col_name = "majors"
     minor_col_name = "minors"
@@ -243,7 +244,6 @@ def heterogeneous_uniform_neighbor_sample(
             " of the libcugraph C++ API"
         )
 
-
     if with_edge_properties:
         warning_msg = (
             "The with_edge_properties flag is deprecated"
@@ -260,8 +260,7 @@ def heterogeneous_uniform_neighbor_sample(
             start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
         )
 
-
-    """ 
+    """
     # No batch_ids, the rank owning the vertices will wom the final
     # result.
     if with_edge_properties and not with_batch_ids:
@@ -292,11 +291,9 @@ def heterogeneous_uniform_neighbor_sample(
 
     start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
 
-
     if G.renumbered:
         start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
 
-
     sampling_result_array_dict = pylibcugraph_heterogeneous_uniform_neighbor_sample(
         resource_handle=ResourceHandle(),
         input_graph=G._plc_graph,
@@ -326,7 +323,7 @@ def heterogeneous_uniform_neighbor_sample(
         return_offsets=return_offsets,
         renumber=renumber,
         use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column, # Deprecated flag
+        include_hop_column=include_hop_column,  # Deprecated flag
     )
 
     if G.renumbered and not renumber:
diff --git a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
index 62ea1a93686..804bea9336d 100644
--- a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
@@ -14,8 +14,9 @@
 from __future__ import annotations
 
 from pylibcugraph import ResourceHandle
-from pylibcugraph import homogeneous_biased_neighbor_sample as \
-    pylibcugraph_homogeneous_biased_neighbor_sample
+from pylibcugraph import (
+    homogeneous_biased_neighbor_sample as pylibcugraph_homogeneous_biased_neighbor_sample,
+)
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
 
@@ -87,7 +88,7 @@ def homogeneous_biased_neighbor_sample(
 
     start_vertex_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
@@ -102,7 +103,7 @@ def homogeneous_biased_neighbor_sample(
         Deprecated.
         Flag to specify whether to return edge properties (weight, edge id,
         edge type, batch id, hop id) with the sampled edges.
-    
+
     prior_sources_behavior: str, optional (default=None)
         Options are "carryover", and "exclude".
         Default will leave the source list as-is.
@@ -110,37 +111,37 @@ def homogeneous_biased_neighbor_sample(
         current hop.
         Exclude will exclude sources from previous hops from reappearing
         as sources in future hops.
-    
+
     deduplicate_sources: bool, optional (default=False)
         Whether to first deduplicate the list of possible sources
         from the previous destinations before performing next
         hop.
-    
+
     return_hops: bool, optional (default=True)
         Whether to return the sampling results with hop ids
         corresponding to the hop where the edge appeared.
         Defaults to True.
-    
+
     renumber: bool, optional (default=False)
         Whether to renumber on a per-batch basis.  If True,
         will return the renumber map and renumber map offsets
         as an additional dataframe.
-    
+
     retain_seeds: bool, optional (default=False)
         If True, will retain the original seeds (original source vertices)
         in the output even if they do not have outgoing neighbors.
-    
+
     compression: str, optional (default=COO)
         Sets the compression type for the output minibatches.
         Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-    
+
     compress_per_hop: bool, optional (default=False)
         Whether to compress globally (default), or to produce a separate
         compressed edgelist per hop.
 
     random_state: int, optional
         Random seed to use when making sampling calls.
-    
+
     return_offsets: bool, optional (default=False)
         Whether to return the sampling results with batch ids
         included as one dataframe, or to instead return two
@@ -211,8 +212,8 @@ def homogeneous_biased_neighbor_sample(
                         Contains the batch offsets for the renumber maps
     """
 
-    use_legacy_names = False # Deprecated parameter
-    include_hop_column=not return_offsets # Deprecated parameter
+    use_legacy_names = False  # Deprecated parameter
+    include_hop_column = not return_offsets  # Deprecated parameter
 
     major_col_name = "majors"
     minor_col_name = "minors"
@@ -238,7 +239,6 @@ def homogeneous_biased_neighbor_sample(
             " of the libcugraph C++ API"
         )
 
-
     if with_edge_properties:
         warning_msg = (
             "The with_edge_properties flag is deprecated"
@@ -255,8 +255,7 @@ def homogeneous_biased_neighbor_sample(
             start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
         )
 
-
-    """ 
+    """
     # No batch_ids, the rank owning the vertices will wom the final
     # result.
     if with_edge_properties and not with_batch_ids:
@@ -287,11 +286,9 @@ def homogeneous_biased_neighbor_sample(
 
     start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
 
-
     if G.renumbered:
         start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
 
-
     sampling_result_array_dict = pylibcugraph_homogeneous_biased_neighbor_sample(
         resource_handle=ResourceHandle(),
         input_graph=G._plc_graph,
@@ -320,7 +317,7 @@ def homogeneous_biased_neighbor_sample(
         return_offsets=return_offsets,
         renumber=renumber,
         use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column, # Deprecated flag
+        include_hop_column=include_hop_column,  # Deprecated flag
     )
 
     if G.renumbered and not renumber:
diff --git a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
index 4c015251af7..6ff955215ca 100644
--- a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
@@ -14,8 +14,9 @@
 from __future__ import annotations
 
 from pylibcugraph import ResourceHandle
-from pylibcugraph import homogeneous_uniform_neighbor_sample as \
-    pylibcugraph_homogeneous_uniform_neighbor_sample
+from pylibcugraph import (
+    homogeneous_uniform_neighbor_sample as pylibcugraph_homogeneous_uniform_neighbor_sample,
+)
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
 
@@ -86,7 +87,7 @@ def homogeneous_uniform_neighbor_sample(
 
     start_vertex_list : list or cudf.Series (int32)
         a list of starting vertices for sampling
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
@@ -101,7 +102,7 @@ def homogeneous_uniform_neighbor_sample(
         Deprecated.
         Flag to specify whether to return edge properties (weight, edge id,
         edge type, batch id, hop id) with the sampled edges.
-    
+
     prior_sources_behavior: str, optional (default=None)
         Options are "carryover", and "exclude".
         Default will leave the source list as-is.
@@ -109,37 +110,37 @@ def homogeneous_uniform_neighbor_sample(
         current hop.
         Exclude will exclude sources from previous hops from reappearing
         as sources in future hops.
-    
+
     deduplicate_sources: bool, optional (default=False)
         Whether to first deduplicate the list of possible sources
         from the previous destinations before performing next
         hop.
-    
+
     return_hops: bool, optional (default=True)
         Whether to return the sampling results with hop ids
         corresponding to the hop where the edge appeared.
         Defaults to True.
-    
+
     renumber: bool, optional (default=False)
         Whether to renumber on a per-batch basis.  If True,
         will return the renumber map and renumber map offsets
         as an additional dataframe.
-    
+
     retain_seeds: bool, optional (default=False)
         If True, will retain the original seeds (original source vertices)
         in the output even if they do not have outgoing neighbors.
-    
+
     compression: str, optional (default=COO)
         Sets the compression type for the output minibatches.
         Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-    
+
     compress_per_hop: bool, optional (default=False)
         Whether to compress globally (default), or to produce a separate
         compressed edgelist per hop.
 
     random_state: int, optional
         Random seed to use when making sampling calls.
-    
+
     return_offsets: bool, optional (default=False)
         Whether to return the sampling results with batch ids
         included as one dataframe, or to instead return two
@@ -210,8 +211,8 @@ def homogeneous_uniform_neighbor_sample(
                         Contains the batch offsets for the renumber maps
     """
 
-    use_legacy_names = False # Deprecated parameter
-    include_hop_column=not return_offsets # Deprecated parameter
+    use_legacy_names = False  # Deprecated parameter
+    include_hop_column = not return_offsets  # Deprecated parameter
 
     major_col_name = "majors"
     minor_col_name = "minors"
@@ -237,7 +238,6 @@ def homogeneous_uniform_neighbor_sample(
             " of the libcugraph C++ API"
         )
 
-
     if with_edge_properties:
         warning_msg = (
             "The with_edge_properties flag is deprecated"
@@ -254,8 +254,7 @@ def homogeneous_uniform_neighbor_sample(
             start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
         )
 
-
-    """ 
+    """
     # No batch_ids, the rank owning the vertices will wom the final
     # result.
     if with_edge_properties and not with_batch_ids:
@@ -286,11 +285,9 @@ def homogeneous_uniform_neighbor_sample(
 
     start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
 
-
     if G.renumbered:
         start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
 
-
     sampling_result_array_dict = pylibcugraph_homogeneous_uniform_neighbor_sample(
         resource_handle=ResourceHandle(),
         input_graph=G._plc_graph,
@@ -319,7 +316,7 @@ def homogeneous_uniform_neighbor_sample(
         return_offsets=return_offsets,
         renumber=renumber,
         use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column, # Deprecated flag
+        include_hop_column=include_hop_column,  # Deprecated flag
     )
 
     if G.renumbered and not renumber:
diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py
index ea584569d77..97c8ab342e6 100644
--- a/python/pylibcugraph/pylibcugraph/__init__.py
+++ b/python/pylibcugraph/pylibcugraph/__init__.py
@@ -40,15 +40,23 @@
 from pylibcugraph.bfs import bfs
 
 from pylibcugraph.uniform_neighbor_sample import uniform_neighbor_sample
-from pylibcugraph.homogeneous_uniform_neighbor_sample import homogeneous_uniform_neighbor_sample
-from pylibcugraph.homogeneous_biased_neighbor_sample import homogeneous_biased_neighbor_sample
-from pylibcugraph.heterogeneous_uniform_neighbor_sample import heterogeneous_uniform_neighbor_sample
-from pylibcugraph.heterogeneous_biased_neighbor_sample import heterogeneous_biased_neighbor_sample
+from pylibcugraph.homogeneous_uniform_neighbor_sample import (
+    homogeneous_uniform_neighbor_sample,
+)
+from pylibcugraph.homogeneous_biased_neighbor_sample import (
+    homogeneous_biased_neighbor_sample,
+)
+from pylibcugraph.heterogeneous_uniform_neighbor_sample import (
+    heterogeneous_uniform_neighbor_sample,
+)
+from pylibcugraph.heterogeneous_biased_neighbor_sample import (
+    heterogeneous_biased_neighbor_sample,
+)
 
 
-#FIXME: break down the API into homogeneous and heterogeneous neighbor sample
-#from pylibcugraph.uniform_neighbor_sample import neighbor_sample
-#from pylibcugraph.biased_neighbor_sample import biased_neighbor_sample
+# FIXME: break down the API into homogeneous and heterogeneous neighbor sample
+# from pylibcugraph.uniform_neighbor_sample import neighbor_sample
+# from pylibcugraph.biased_neighbor_sample import biased_neighbor_sample
 
 from pylibcugraph.negative_sampling import negative_sampling
 
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
index 401b2ee6258..ea71caaa7db 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
@@ -113,23 +113,23 @@ def heterogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
 
     input_graph : SGGraph or MGGraph
         The input graph, for either Single or Multi-GPU operations.
-    
+
     edge_biases: FIXE: update this - Create edge_biases of type
         'cugraph_edge_property_view_t' - edge biases not yet supported.
 
     start_vertex_list: device array type
         Device array containing the list of starting vertices for sampling.
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     h_fan_out: tuple of numpy array type
         Device array containing the branching out (fan-out) degrees per
         starting vertex for each hop level
-        
+
         The sampling method can use different fan_out values for each edge type
         which is not the case for homogeneous neighborhood sampling.
-    
+
     num_edge_types: Number of edge types where a value of 1 translates to homogeneous neighbor
         sample whereas a value greater than 1 translates to heterogeneous neighbor
         sample.
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
index 477a010037b..2580c7369be 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
@@ -113,17 +113,17 @@ def heterogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
 
     start_vertex_list: device array type
         Device array containing the list of starting vertices for sampling.
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     h_fan_out: tuple of numpy array type
         Device array containing the branching out (fan-out) degrees per
         starting vertex for each hop level
-        
+
         The sampling method can use different fan_out values for each edge type
         which is not the case for homogeneous neighborhood sampling.
-    
+
     num_edge_types: Number of edge types where a value of 1 translates to homogeneous neighbor
         sample whereas a value greater than 1 translates to heterogeneous neighbor
         sample.
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
index d95d8a80a2e..51b770f0fab 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
@@ -112,20 +112,20 @@ def homogeneous_biased_neighbor_sample(ResourceHandle resource_handle,
 
     input_graph : SGGraph or MGGraph
         The input graph, for either Single or Multi-GPU operations.
-    
+
     edge_biases: FIXE: update this - Create edge_biases of type
         'cugraph_edge_property_view_t' - edge biases not yet supported.
 
     start_vertex_list: device array type
         Device array containing the list of starting vertices for sampling.
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     h_fan_out: tuple of numpy array type
         Device array containing the branching out (fan-out) degrees per
         starting vertex for each hop level
-        
+
         The sampling method can use different fan_out values for each edge type
         which is not the case for homogeneous neighborhood sampling.
 
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
index 1f5ec94705a..bddcc0b29ff 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
@@ -112,14 +112,14 @@ def homogeneous_uniform_neighbor_sample(ResourceHandle resource_handle,
 
     start_vertex_list: device array type
         Device array containing the list of starting vertices for sampling.
-    
+
     start_vertex_offsets: list[int] (Optional)
         Offsets of each label within the start vertex list.
 
     h_fan_out: tuple of numpy array type
         Device array containing the branching out (fan-out) degrees per
         starting vertex for each hop level
-        
+
         The sampling method can use different fan_out values for each edge type
         which is not the case for homogeneous neighborhood sampling.
 

From 6b547fb83f140d8d1347b6c0d80406a0a49dea7d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 12:16:04 -0700
Subject: [PATCH 084/117] fix style

---
 .../cugraph/sampling/heterogeneous_biased_neighbor_sample.py   | 3 ++-
 .../cugraph/sampling/heterogeneous_uniform_neighbor_sample.py  | 3 ++-
 .../sampling/homogeneous_biased_uniform_neighbor_sample.py     | 3 ++-
 .../cugraph/sampling/homogeneous_uniform_neighbor_sample.py    | 3 ++-
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
index 0266bb6ddb6..380e7d371ad 100644
--- a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
@@ -15,7 +15,8 @@
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import (
-    heterogeneous_biased_neighbor_sample as pylibcugraph_heterogeneous_biased_neighbor_sample,
+    heterogeneous_biased_neighbor_sample as \
+        pylibcugraph_heterogeneous_biased_neighbor_sample,
 )
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
diff --git a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
index 3c81877abc6..b0c1ccdc9ce 100644
--- a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
@@ -15,7 +15,8 @@
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import (
-    heterogeneous_uniform_neighbor_sample as pylibcugraph_heterogeneous_uniform_neighbor_sample,
+    heterogeneous_uniform_neighbor_sample as \
+        pylibcugraph_heterogeneous_uniform_neighbor_sample,
 )
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
diff --git a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
index 804bea9336d..27244ba0f96 100644
--- a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
@@ -15,7 +15,8 @@
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import (
-    homogeneous_biased_neighbor_sample as pylibcugraph_homogeneous_biased_neighbor_sample,
+    homogeneous_biased_neighbor_sample as \
+        pylibcugraph_homogeneous_biased_neighbor_sample,
 )
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
diff --git a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
index 6ff955215ca..b980e38407c 100644
--- a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
@@ -15,7 +15,8 @@
 
 from pylibcugraph import ResourceHandle
 from pylibcugraph import (
-    homogeneous_uniform_neighbor_sample as pylibcugraph_homogeneous_uniform_neighbor_sample,
+    homogeneous_uniform_neighbor_sample as \
+        pylibcugraph_homogeneous_uniform_neighbor_sample,
 )
 
 from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict

From 14e9a9931be14647ca4541d7d5cf4acc6f8ec700 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 16:11:11 -0700
Subject: [PATCH 085/117] update cmakelist

---
 cpp/CMakeLists.txt | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index f23d9db9b8c..a8d9149ca2e 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -282,10 +282,10 @@ set(CUGRAPH_SOURCES
     src/sampling/detail/sample_edges_mg_v32_e32.cu
     src/sampling/detail/shuffle_and_organize_output_mg_v64_e64.cu
     src/sampling/detail/shuffle_and_organize_output_mg_v32_e32.cu
-    src/sampling/neighbor_sampling_mg_v32_e32.cpp
-    src/sampling/neighbor_sampling_mg_v64_e64.cpp
-    src/sampling/neighbor_sampling_sg_v32_e32.cpp
-    src/sampling/neighbor_sampling_sg_v64_e64.cpp
+    src/sampling/neighbor_sampling_mg_v32_e32.cu
+    src/sampling/neighbor_sampling_mg_v64_e64.cu
+    src/sampling/neighbor_sampling_sg_v32_e32.cu
+    src/sampling/neighbor_sampling_sg_v64_e64.cu
     src/sampling/negative_sampling_sg_v32_e32.cu
     src/sampling/negative_sampling_sg_v64_e64.cu
     src/sampling/negative_sampling_mg_v32_e32.cu

From aebfd086ec67a8e76f640d7aec40bfdcb244f955 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 22 Oct 2024 16:12:00 -0700
Subject: [PATCH 086/117] update type combination

---
 .../shuffle_and_organize_output_mg_v32_e64.cu | 61 +++++++++++++++++++
 ...heterogeneous_biased_neighbor_sampling.cpp | 14 -----
 ...eterogeneous_uniform_neighbor_sampling.cpp | 16 +----
 .../homogeneous_biased_neighbor_sampling.cpp  | 12 ----
 .../homogeneous_uniform_neighbor_sampling.cpp | 17 ++----
 ...heterogeneous_biased_neighbor_sampling.cpp |  6 --
 ...eterogeneous_uniform_neighbor_sampling.cpp |  6 --
 ...g_homogeneous_biased_neighbor_sampling.cpp |  6 --
 ..._homogeneous_uniform_neighbor_sampling.cpp |  6 --
 9 files changed, 68 insertions(+), 76 deletions(-)
 create mode 100644 cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu

diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
new file mode 100644
index 00000000000..ef760844b51
--- /dev/null
+++ b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "sampling/detail/shuffle_and_organize_output_impl.cuh"
+
+namespace cugraph {
+namespace detail {
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<float>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+shuffle_and_organize_output(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& majors,
+  rmm::device_uvector<int32_t>&& minors,
+  std::optional<rmm::device_uvector<float>>&& weights,
+  std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+  std::optional<rmm::device_uvector<int32_t>>&& edge_types,
+  std::optional<rmm::device_uvector<int32_t>>&& hops,
+  std::optional<rmm::device_uvector<int32_t>>&& labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank);
+
+template std::tuple<rmm::device_uvector<int32_t>,
+                    rmm::device_uvector<int32_t>,
+                    std::optional<rmm::device_uvector<double>>,
+                    std::optional<rmm::device_uvector<int64_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<int32_t>>,
+                    std::optional<rmm::device_uvector<size_t>>>
+shuffle_and_organize_output(
+  raft::handle_t const& handle,
+  rmm::device_uvector<int32_t>&& majors,
+  rmm::device_uvector<int32_t>&& minors,
+  std::optional<rmm::device_uvector<double>>&& weights,
+  std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
+  std::optional<rmm::device_uvector<int32_t>>&& edge_types,
+  std::optional<rmm::device_uvector<int32_t>>&& hops,
+  std::optional<rmm::device_uvector<int32_t>>&& labels,
+  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank);
+
+}  // namespace detail
+}  // namespace cugraph
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index c8bfb654598..3c3739ed668 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -75,8 +75,6 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
     auto edge_weight_view =
       edge_weights ? std::make_optional((*edge_weights).view()) : std::nullopt;
 
-    std::optional<cugraph::edge_property_t<decltype(graph_view), bool>> edge_mask{std::nullopt};
-
     constexpr float select_probability{0.05};
 
     // FIXME:  Update the tests to initialize RngState and use it instead
@@ -251,12 +249,6 @@ TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_File, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
@@ -269,12 +261,6 @@ TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 15be79a623e..83f99e68e9a 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -116,11 +116,11 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
     // Generate the edge types
 
-    std::optional<cugraph::edge_property_t<decltype(graph_view), edge_type_t>> edge_types{
+    std::optional<cugraph::edge_property_t<decltype(graph_view), int32_t>> edge_types{
       std::nullopt};
 
     if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
-      edge_types = cugraph::test::generate<decltype(graph_view), edge_type_t>::edge_property(
+      edge_types = cugraph::test::generate<decltype(graph_view), int32_t>::edge_property(
         handle, graph_view, heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
     }
 
@@ -249,12 +249,6 @@ TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_File, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
@@ -267,12 +261,6 @@ TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
index d9f3e90f127..d01206c5cef 100644
--- a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -233,12 +233,6 @@ TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_File, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
@@ -251,12 +245,6 @@ TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Homogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
index 8c50ab56917..508914e2495 100644
--- a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -80,6 +80,11 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
       graph_view.attach_edge_mask((*edge_mask).view());
     }
 
+    // FIXME: Read a tuple of two edge mask and mask out if edge mask is set in either 1 (OR) and create
+    // a new one.
+    // No graph view can have two mask and perform OR in itself, and need to OR the mask 
+    // manually by itself.
+
     constexpr float select_probability{0.05};
 
     // FIXME:  Update the tests to initialize RngState and use it instead
@@ -231,12 +236,6 @@ TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int32Float)
     override_File_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_File_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_File, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
@@ -249,12 +248,6 @@ TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_Homogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index c2e34355062..688d8480118 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -338,12 +338,6 @@ TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index 6282f57f1be..18e75d1b1b6 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -339,12 +339,6 @@ TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Floa
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
index 4ad5f0a2053..ec134938067 100644
--- a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
@@ -323,12 +323,6 @@ TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_MGHomogeneous_Biased_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(
diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
index 68ca7d4dd8e..e33000044a5 100644
--- a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -323,12 +323,6 @@ TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int32Float)
     override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
 }
 
-TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt32Int64Float)
-{
-  run_current_test<int32_t, int64_t, float>(
-    override_Rmat_Usecase_with_cmd_line_arguments(GetParam()));
-}
-
 TEST_P(Tests_MGHomogeneous_Uniform_Neighbor_Sampling_Rmat, CheckInt64Int64Float)
 {
   run_current_test<int64_t, int64_t, float>(

From b159e1b7909d251ace4694be0fe96a8b65cc106d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 29 Oct 2024 20:34:58 -0700
Subject: [PATCH 087/117] fix symbol lookup error

---
 cpp/src/c_api/neighbor_sampling.cpp | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 5b05b542015..cc4c09ead2f 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -875,19 +875,20 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           handle_.get_comms(), num_local_labels, handle_.get_stream());
 
         std::exclusive_scan(
-          global_labels.begin(), global_labels.end(), global_labels.begin(), size_t{0});
-
-        // Compute the global start_vertex_label_offsets
-        cugraph::detail::transform_increment(handle_.get_stream(),
-                                             start_vertex_offsets_->as_type<size_t>(),
-                                             start_vertex_offsets_->size_,
-                                             global_labels[handle_.get_comms().get_rank()]);
+          global_labels.begin(), global_labels.end(), global_labels.begin(), label_t{0});
 
         // Retrieve the start_vertex_labels
         start_vertex_labels = cugraph::detail::convert_starting_vertex_offsets_to_labels(
           handle_,
           raft::device_span<size_t const>{start_vertex_offsets_->as_type<size_t>(),
                                           start_vertex_offsets_->size_});
+
+        // Compute the global start_vertex_label_offsets
+        cugraph::detail::transform_increment(handle_.get_stream(),
+                                              (label_t*)(*start_vertex_labels).data(),
+                                              (size_t)(*start_vertex_labels).size(),
+                                              (label_t)global_labels[handle_.get_comms().get_rank()]
+                                              );
       }
 
       if constexpr (multi_gpu) {

From 3b0c016aa0886962af773714b09eb506a19c76bc Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 08:44:22 -0700
Subject: [PATCH 088/117] leverage raft span instead of raw pointers

---
 .../cugraph/detail/utility_wrappers.hpp       |  9 ++--
 cpp/src/c_api/neighbor_sampling.cpp           |  9 ++--
 cpp/src/detail/utility_wrappers_32.cu         | 15 +++---
 cpp/src/detail/utility_wrappers_64.cu         | 15 +++---
 cpp/src/detail/utility_wrappers_impl.cuh      | 47 ++++++++++---------
 5 files changed, 48 insertions(+), 47 deletions(-)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index a7400a43039..a9486621798 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -73,11 +73,10 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
  * @param[out]  d_value      device array to sort
- * @param[in]   size         number of elements in array
  *
  */
 template <typename value_t>
-void sort(raft::handle_t const& handle, value_t* d_value, size_t size);
+void sort(raft::handle_t const& handle, raft::device_span<value_t> d_value);
 
 /**
  * @brief    Keep unique element from a buffer
@@ -87,12 +86,11 @@ void sort(raft::handle_t const& handle, value_t* d_value, size_t size);
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
  * @param[in]  d_value      device array of unique elements.
- * @param[in]   size         number of elements in array
  * @return the number of unique elements.
  *
  */
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size);
+size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value);
 
 /**
  * @brief    Increment the values of a buffer by a constant value
@@ -102,12 +100,11 @@ size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size);
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
  * @param[out]  d_value      device array to update
- * @param[in]   size         number of elements in array
  * @param[in]   value        value to be added to each element of the buffer
  *
  */
 template <typename value_t>
-void transform_increment(raft::handle_t const& handle, value_t* d_value, size_t size, size_t value);
+void transform_increment(rmm::cuda_stream_view const& stream_view, raft::device_span<value_t> d_value, value_t value);
 
 /**
  * @brief    Fill a buffer with a sequence of values
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index cc4c09ead2f..527fa53ff65 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -884,9 +884,9 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
                                           start_vertex_offsets_->size_});
 
         // Compute the global start_vertex_label_offsets
+
         cugraph::detail::transform_increment(handle_.get_stream(),
-                                              (label_t*)(*start_vertex_labels).data(),
-                                              (size_t)(*start_vertex_labels).size(),
+                                              raft::device_span<label_t>{(*start_vertex_labels).data(), (*start_vertex_labels).size()},
                                               (label_t)global_labels[handle_.get_comms().get_rank()]
                                               );
       }
@@ -902,9 +902,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
           // Get unique labels
           // sort the start_vertex_labels
-          cugraph::detail::sort(handle_.get_stream(), unique_labels.begin(), unique_labels.size());
+          cugraph::detail::sort(handle_.get_stream(), raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
+          
           auto num_unique_labels = cugraph::detail::unique(
-            handle_.get_stream(), unique_labels.begin(), unique_labels.size());
+            handle_.get_stream(), raft::device_span<label_t const>{unique_labels.data(), unique_labels.size()});
 
           (*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 
diff --git a/cpp/src/detail/utility_wrappers_32.cu b/cpp/src/detail/utility_wrappers_32.cu
index 67e86b21213..3fe1a507870 100644
--- a/cpp/src/detail/utility_wrappers_32.cu
+++ b/cpp/src/detail/utility_wrappers_32.cu
@@ -63,9 +63,10 @@ template void scalar_fill(raft::handle_t const& handle, size_t* d_value, size_t
 
 template void scalar_fill(raft::handle_t const& handle, float* d_value, size_t size, float value);
 
-template void sort(raft::handle_t const& handle, int32_t* d_value, size_t size);
+template void sort(raft::handle_t const& handle, raft::device_span<int32_t> d_value);
 
-template size_t unique(raft::handle_t const& handle, int32_t* d_value, size_t size);
+template size_t unique(raft::handle_t const& handle, raft::device_span<int32_t> d_value);
+template size_t unique(raft::handle_t const& handle, raft::device_span<uint32_t> d_value);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int32_t* d_value,
@@ -78,14 +79,12 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             uint32_t start_value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  int32_t* d_value,
-                                  size_t size,
-                                  size_t value);
+                                  raft::device_span<int32_t> d_value,
+                                  int32_t value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  uint32_t* d_value,
-                                  size_t size,
-                                  size_t value);
+                                  raft::device_span<uint32_t> d_value,
+                                  uint32_t value);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int32_t* d_value,
diff --git a/cpp/src/detail/utility_wrappers_64.cu b/cpp/src/detail/utility_wrappers_64.cu
index e75fdf7b075..4012aa69b5c 100644
--- a/cpp/src/detail/utility_wrappers_64.cu
+++ b/cpp/src/detail/utility_wrappers_64.cu
@@ -61,9 +61,10 @@ template void scalar_fill(raft::handle_t const& handle,
 
 template void scalar_fill(raft::handle_t const& handle, double* d_value, size_t size, double value);
 
-template void sort(raft::handle_t const& handle, int64_t* d_value, size_t size);
+template void sort(raft::handle_t const& handle, raft::device_span<int64_t> d_value);
 
-template size_t unique(raft::handle_t const& handle, int32_t* d_value, size_t size);
+template size_t unique(raft::handle_t const& handle, raft::device_span<int64_t> d_value);
+template size_t unique(raft::handle_t const& handle, raft::device_span<uint64_t> d_value);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int64_t* d_value,
@@ -76,14 +77,12 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             uint64_t start_value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  int64_t* d_value,
-                                  size_t size,
-                                  size_t value);
+                                  raft::device_span<int64_t> d_value,
+                                  int64_t value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  uint64_t* d_value,
-                                  size_t size,
-                                  size_t value);
+                                  raft::device_span<uint64_t> d_value,
+                                  uint64_t value);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int64_t* d_value,
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index eaca2516dbf..60bb1557019 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -65,19 +65,16 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
 }
 
 template <typename value_t>
-void sort(raft::handle_t const& handle, value_t* d_value, size_t size)
+void sort(raft::handle_t const& handle, raft::device_span<value_t> d_value)
 {
-  thrust::sort(handle.get_thrust_policy(), d_value, d_value + size);
+  thrust::sort(handle.get_thrust_policy(), d_value.begin(), d_value.end());
 }
 
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, value_t* d_value, size_t size)
+size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value)
 {
-  // auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
-  auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value, d_value + size);
-  // auto num_unique_element =
-  return thrust::distance(d_value, unique_element_last);
-  // masked_edgelist_srcs.resize(2* masked_edgelist_srcs.size(), handle.get_stream());
+  auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value.begin(), d_value.end());
+  return thrust::distance(d_value.begin(), unique_element_last);
 }
 
 template <typename value_t>
@@ -89,20 +86,28 @@ void sequence_fill(rmm::cuda_stream_view const& stream_view,
   thrust::sequence(rmm::exec_policy(stream_view), d_value, d_value + size, start_value);
 }
 
+
 template <typename value_t>
-void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         value_t* d_value,
-                         size_t size,
-                         size_t incr)
-{
-  thrust::transform(rmm::exec_policy(stream_view),
-                    d_value,
-                    d_value + size,
-                    d_value,
-                    cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
-                      return static_cast<value_t>(value + incr);
-                    }));
-}
+ void transform_increment(rmm::cuda_stream_view const& stream_view,
+                          raft::device_span<value_t> d_value,
+                          value_t incr)
+ {
+   thrust::transform(rmm::exec_policy(stream_view),
+                     d_value.begin(),
+                     d_value.end(),
+                     d_value.begin(),
+                     cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
+                       return static_cast<value_t>(value + incr);
+                     }));
+ }
+
+
+template <typename value_t>
+ void transform_increment_(rmm::cuda_stream_view const& stream_view,
+                           value_t d_value)
+ {
+
+ }
 
 template <typename value_t>
 void stride_fill(rmm::cuda_stream_view const& stream_view,

From da82567bab80e66fb98de511037056148c06ca22 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 09:50:17 -0700
Subject: [PATCH 089/117] remove python implementation of heterogeneous
 neighborhood sampling and push it to a different PR

---
 .../heterogeneous_biased_neighbor_sample.py   | 337 ------------------
 .../heterogeneous_uniform_neighbor_sample.py  | 337 ------------------
 ...ogeneous_biased_uniform_neighbor_sample.py | 331 -----------------
 .../homogeneous_uniform_neighbor_sample.py    | 330 -----------------
 4 files changed, 1335 deletions(-)
 delete mode 100644 python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
 delete mode 100644 python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
 delete mode 100644 python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
 delete mode 100644 python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py

diff --git a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
deleted file mode 100644
index 380e7d371ad..00000000000
--- a/python/cugraph/cugraph/sampling/heterogeneous_biased_neighbor_sample.py
+++ /dev/null
@@ -1,337 +0,0 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from pylibcugraph import ResourceHandle
-from pylibcugraph import (
-    heterogeneous_biased_neighbor_sample as \
-        pylibcugraph_heterogeneous_biased_neighbor_sample,
-)
-
-from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
-
-import numpy
-
-import cudf
-import cupy as cp
-import warnings
-
-from typing import Union, Tuple, Sequence, List
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from cugraph import Graph
-
-
-start_col_name = "_START_"
-batch_col_name = "_BATCH_"
-
-
-# FIXME: Move this function to the utility module so that it can be
-# shared by other algos
-def ensure_valid_dtype(input_graph, start_vertex_list):
-    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
-    if isinstance(start_vertex_list, cudf.Series):
-        start_vertex_list_dtypes = start_vertex_list.dtype
-    else:
-        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
-
-    if start_vertex_list_dtypes != vertex_dtype:
-        warning_msg = (
-            "biased neighbor sample requires 'start_vertex_list' to match the graph's "
-            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
-            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
-        )
-        warnings.warn(warning_msg, UserWarning)
-        start_vertex_list = start_vertex_list.astype(vertex_dtype)
-
-    return start_vertex_list
-
-
-def heterogeneous_biased_neighbor_sample(
-    G: Graph,
-    start_vertex_list: Sequence,
-    start_vertex_offsets: Sequence,
-    fanout_vals: List[int],
-    num_edge_types: int,
-    with_replacement: bool = True,
-    with_edge_properties: bool = False,  # deprecated
-    prior_sources_behavior: str = None,
-    deduplicate_sources: bool = False,
-    return_hops: bool = True,
-    renumber: bool = False,
-    retain_seeds: bool = False,
-    compression: str = "COO",
-    compress_per_hop: bool = False,
-    random_state: int = None,
-    return_offsets: bool = False,
-) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
-    """
-    Does neighborhood sampling, which samples nodes from a graph based on the
-    current node's neighbors, with a corresponding fanout value at each hop.
-
-    Parameters
-    ----------
-    G : cugraph.Graph
-        cuGraph graph, which contains connectivity information as dask cudf
-        edge list dataframe
-
-    start_vertex_list : list or cudf.Series (int32)
-        a list of starting vertices for sampling
-
-    start_vertex_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
-    fanout_vals : list (int32)
-        List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
-
-    num_edge_types: int32
-        Number of edge types where a value of 1 translates to homogeneous neighbor
-        sample whereas a value greater than 1 translates to heterogeneous neighbor
-        sample.
-
-    with_replacement: bool, optional (default=True)
-        Flag to specify if the random sampling is done with replacement
-
-    with_edge_properties: bool, optional (default=False)
-        Deprecated.
-        Flag to specify whether to return edge properties (weight, edge id,
-        edge type, batch id, hop id) with the sampled edges.
-
-    prior_sources_behavior: str, optional (default=None)
-        Options are "carryover", and "exclude".
-        Default will leave the source list as-is.
-        Carryover will carry over sources from previous hops to the
-        current hop.
-        Exclude will exclude sources from previous hops from reappearing
-        as sources in future hops.
-
-    deduplicate_sources: bool, optional (default=False)
-        Whether to first deduplicate the list of possible sources
-        from the previous destinations before performing next
-        hop.
-
-    return_hops: bool, optional (default=True)
-        Whether to return the sampling results with hop ids
-        corresponding to the hop where the edge appeared.
-        Defaults to True.
-
-    renumber: bool, optional (default=False)
-        Whether to renumber on a per-batch basis.  If True,
-        will return the renumber map and renumber map offsets
-        as an additional dataframe.
-
-    retain_seeds: bool, optional (default=False)
-        If True, will retain the original seeds (original source vertices)
-        in the output even if they do not have outgoing neighbors.
-
-    compression: str, optional (default=COO)
-        Sets the compression type for the output minibatches.
-        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-
-    compress_per_hop: bool, optional (default=False)
-        Whether to compress globally (default), or to produce a separate
-        compressed edgelist per hop.
-
-    random_state: int, optional
-        Random seed to use when making sampling calls.
-
-    return_offsets: bool, optional (default=False)
-        Whether to return the sampling results with batch ids
-        included as one dataframe, or to instead return two
-        dataframes, one with sampling results and one with
-        batch ids and their start offsets.
-
-    Returns
-    -------
-    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
-        GPU data frame containing multiple cudf.Series
-
-        If with_edge_properties=False:
-            df['sources']: cudf.Series
-                Contains the source vertices from the sampling result
-            df['destinations']: cudf.Series
-                Contains the destination vertices from the sampling result
-            df['indices']: cudf.Series
-                Contains the indices (edge weights) from the sampling result
-                for path reconstruction
-
-        If with_edge_properties=True:
-            If return_offsets=False:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-
-            If return_offsets=True:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-
-                offsets_df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                offsets_df['offsets']: cudf.Series
-                    Contains the offsets of each batch in the sampling result
-
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-    """
-
-    use_legacy_names = False  # Deprecated parameter
-    include_hop_column = not return_offsets  # Deprecated parameter
-
-    major_col_name = "majors"
-    minor_col_name = "minors"
-
-    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
-        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
-
-    if (
-        (compression != "COO")
-        and (not compress_per_hop)
-        and prior_sources_behavior != "exclude"
-    ):
-        raise ValueError(
-            "hop-agnostic compression is only supported with"
-            " the exclude prior sources behavior due to limitations "
-            "of the libcugraph C++ API"
-        )
-
-    if compress_per_hop and prior_sources_behavior != "carryover":
-        raise ValueError(
-            "Compressing the edgelist per hop is only supported "
-            "with the carryover prior sources behavior due to limitations"
-            " of the libcugraph C++ API"
-        )
-
-    if with_edge_properties:
-        warning_msg = (
-            "The with_edge_properties flag is deprecated"
-            " and will be removed in the next release in favor"
-            " of returning all properties in the graph"
-        )
-        warnings.warn(warning_msg, FutureWarning)
-
-    if isinstance(start_vertex_list, int):
-        start_vertex_list = [start_vertex_list]
-
-    if isinstance(start_vertex_list, list):
-        start_vertex_list = cudf.Series(
-            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
-        )
-
-    """
-    # No batch_ids, the rank owning the vertices will wom the final
-    # result.
-    if with_edge_properties and not with_batch_ids:
-        if isinstance(start_vertex_list, cudf.Series):
-            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
-
-        start_vertex_list[batch_col_name] = cudf.Series(
-            cp.zeros(len(start_vertex_list), dtype="int32")
-        )
-    """
-
-    # fanout_vals must be passed to pylibcugraph as a host array
-    if isinstance(fanout_vals, numpy.ndarray):
-        fanout_vals = fanout_vals.astype("int32")
-    elif isinstance(fanout_vals, list):
-        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
-    elif isinstance(fanout_vals, cp.ndarray):
-        fanout_vals = fanout_vals.get().astype("int32")
-    elif isinstance(fanout_vals, cudf.Series):
-        fanout_vals = fanout_vals.values_host.astype("int32")
-    else:
-        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
-
-    if "weights" in G.edgelist.edgelist_df:
-        weight_t = G.edgelist.edgelist_df["weights"].dtype
-    else:
-        weight_t = "float32"
-
-    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
-
-    if G.renumbered:
-        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
-
-    sampling_result_array_dict = pylibcugraph_heterogeneous_biased_neighbor_sample(
-        resource_handle=ResourceHandle(),
-        input_graph=G._plc_graph,
-        start_vertex_list=start_vertex_list,
-        start_vertex_offsets=start_vertex_offsets,
-        h_fan_out=fanout_vals,
-        num_edge_types=num_edge_types,
-        with_replacement=with_replacement,
-        do_expensive_check=False,
-        with_edge_properties=with_edge_properties,
-        prior_sources_behavior=prior_sources_behavior,
-        deduplicate_sources=deduplicate_sources,
-        return_hops=return_hops,
-        renumber=renumber,
-        retain_seeds=retain_seeds,
-        compression=compression,
-        compress_per_hop=compress_per_hop,
-        random_state=random_state,
-        return_dict=True,
-    )
-
-    dfs = sampling_results_from_cupy_array_dict(
-        sampling_result_array_dict,
-        weight_t,
-        len(fanout_vals),
-        with_edge_properties=with_edge_properties,
-        return_offsets=return_offsets,
-        renumber=renumber,
-        use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column,  # Deprecated flag
-    )
-
-    if G.renumbered and not renumber:
-        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
-        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
-
-    if len(dfs) > 1:
-        return dfs
-
-    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
deleted file mode 100644
index b0c1ccdc9ce..00000000000
--- a/python/cugraph/cugraph/sampling/heterogeneous_uniform_neighbor_sample.py
+++ /dev/null
@@ -1,337 +0,0 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from pylibcugraph import ResourceHandle
-from pylibcugraph import (
-    heterogeneous_uniform_neighbor_sample as \
-        pylibcugraph_heterogeneous_uniform_neighbor_sample,
-)
-
-from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
-
-import numpy
-
-import cudf
-import cupy as cp
-import warnings
-
-from typing import Union, Tuple, Sequence, List
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from cugraph import Graph
-
-
-start_col_name = "_START_"
-batch_col_name = "_BATCH_"
-
-
-# FIXME: Move this function to the utility module so that it can be
-# shared by other algos
-def ensure_valid_dtype(input_graph, start_vertex_list):
-    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
-    if isinstance(start_vertex_list, cudf.Series):
-        start_vertex_list_dtypes = start_vertex_list.dtype
-    else:
-        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
-
-    if start_vertex_list_dtypes != vertex_dtype:
-        warning_msg = (
-            "Uniform neighbor sample requires 'start_vertex_list' to match the graph's "
-            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
-            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
-        )
-        warnings.warn(warning_msg, UserWarning)
-        start_vertex_list = start_vertex_list.astype(vertex_dtype)
-
-    return start_vertex_list
-
-
-def heterogeneous_uniform_neighbor_sample(
-    G: Graph,
-    start_vertex_list: Sequence,
-    start_vertex_offsets: Sequence,
-    fanout_vals: List[int],
-    num_edge_types: int,
-    with_replacement: bool = True,
-    with_edge_properties: bool = False,  # deprecated
-    prior_sources_behavior: str = None,
-    deduplicate_sources: bool = False,
-    return_hops: bool = True,
-    renumber: bool = False,
-    retain_seeds: bool = False,
-    compression: str = "COO",
-    compress_per_hop: bool = False,
-    random_state: int = None,
-    return_offsets: bool = False,
-) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
-    """
-    Does neighborhood sampling, which samples nodes from a graph based on the
-    current node's neighbors, with a corresponding fanout value at each hop.
-
-    Parameters
-    ----------
-    G : cugraph.Graph
-        cuGraph graph, which contains connectivity information as dask cudf
-        edge list dataframe
-
-    start_vertex_list : list or cudf.Series (int32)
-        a list of starting vertices for sampling
-
-    start_vertex_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
-    fanout_vals : list (int32)
-        List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
-
-    num_edge_types: int32
-        Number of edge types where a value of 1 translates to homogeneous neighbor
-        sample whereas a value greater than 1 translates to heterogeneous neighbor
-        sample.
-
-    with_replacement: bool, optional (default=True)
-        Flag to specify if the random sampling is done with replacement
-
-    with_edge_properties: bool, optional (default=False)
-        Deprecated.
-        Flag to specify whether to return edge properties (weight, edge id,
-        edge type, batch id, hop id) with the sampled edges.
-
-    prior_sources_behavior: str, optional (default=None)
-        Options are "carryover", and "exclude".
-        Default will leave the source list as-is.
-        Carryover will carry over sources from previous hops to the
-        current hop.
-        Exclude will exclude sources from previous hops from reappearing
-        as sources in future hops.
-
-    deduplicate_sources: bool, optional (default=False)
-        Whether to first deduplicate the list of possible sources
-        from the previous destinations before performing next
-        hop.
-
-    return_hops: bool, optional (default=True)
-        Whether to return the sampling results with hop ids
-        corresponding to the hop where the edge appeared.
-        Defaults to True.
-
-    renumber: bool, optional (default=False)
-        Whether to renumber on a per-batch basis.  If True,
-        will return the renumber map and renumber map offsets
-        as an additional dataframe.
-
-    retain_seeds: bool, optional (default=False)
-        If True, will retain the original seeds (original source vertices)
-        in the output even if they do not have outgoing neighbors.
-
-    compression: str, optional (default=COO)
-        Sets the compression type for the output minibatches.
-        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-
-    compress_per_hop: bool, optional (default=False)
-        Whether to compress globally (default), or to produce a separate
-        compressed edgelist per hop.
-
-    random_state: int, optional
-        Random seed to use when making sampling calls.
-
-    return_offsets: bool, optional (default=False)
-        Whether to return the sampling results with batch ids
-        included as one dataframe, or to instead return two
-        dataframes, one with sampling results and one with
-        batch ids and their start offsets.
-
-    Returns
-    -------
-    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
-        GPU data frame containing multiple cudf.Series
-
-        If with_edge_properties=False:
-            df['sources']: cudf.Series
-                Contains the source vertices from the sampling result
-            df['destinations']: cudf.Series
-                Contains the destination vertices from the sampling result
-            df['indices']: cudf.Series
-                Contains the indices (edge weights) from the sampling result
-                for path reconstruction
-
-        If with_edge_properties=True:
-            If return_offsets=False:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-
-            If return_offsets=True:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-
-                offsets_df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                offsets_df['offsets']: cudf.Series
-                    Contains the offsets of each batch in the sampling result
-
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-    """
-
-    use_legacy_names = False  # Deprecated parameter
-    include_hop_column = not return_offsets  # Deprecated parameter
-
-    major_col_name = "majors"
-    minor_col_name = "minors"
-
-    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
-        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
-
-    if (
-        (compression != "COO")
-        and (not compress_per_hop)
-        and prior_sources_behavior != "exclude"
-    ):
-        raise ValueError(
-            "hop-agnostic compression is only supported with"
-            " the exclude prior sources behavior due to limitations "
-            "of the libcugraph C++ API"
-        )
-
-    if compress_per_hop and prior_sources_behavior != "carryover":
-        raise ValueError(
-            "Compressing the edgelist per hop is only supported "
-            "with the carryover prior sources behavior due to limitations"
-            " of the libcugraph C++ API"
-        )
-
-    if with_edge_properties:
-        warning_msg = (
-            "The with_edge_properties flag is deprecated"
-            " and will be removed in the next release in favor"
-            " of returning all properties in the graph"
-        )
-        warnings.warn(warning_msg, FutureWarning)
-
-    if isinstance(start_vertex_list, int):
-        start_vertex_list = [start_vertex_list]
-
-    if isinstance(start_vertex_list, list):
-        start_vertex_list = cudf.Series(
-            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
-        )
-
-    """
-    # No batch_ids, the rank owning the vertices will wom the final
-    # result.
-    if with_edge_properties and not with_batch_ids:
-        if isinstance(start_vertex_list, cudf.Series):
-            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
-
-        start_vertex_list[batch_col_name] = cudf.Series(
-            cp.zeros(len(start_vertex_list), dtype="int32")
-        )
-    """
-
-    # fanout_vals must be passed to pylibcugraph as a host array
-    if isinstance(fanout_vals, numpy.ndarray):
-        fanout_vals = fanout_vals.astype("int32")
-    elif isinstance(fanout_vals, list):
-        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
-    elif isinstance(fanout_vals, cp.ndarray):
-        fanout_vals = fanout_vals.get().astype("int32")
-    elif isinstance(fanout_vals, cudf.Series):
-        fanout_vals = fanout_vals.values_host.astype("int32")
-    else:
-        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
-
-    if "weights" in G.edgelist.edgelist_df:
-        weight_t = G.edgelist.edgelist_df["weights"].dtype
-    else:
-        weight_t = "float32"
-
-    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
-
-    if G.renumbered:
-        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
-
-    sampling_result_array_dict = pylibcugraph_heterogeneous_uniform_neighbor_sample(
-        resource_handle=ResourceHandle(),
-        input_graph=G._plc_graph,
-        start_vertex_list=start_vertex_list,
-        start_vertex_offsets=start_vertex_offsets,
-        h_fan_out=fanout_vals,
-        num_edge_types=num_edge_types,
-        with_replacement=with_replacement,
-        do_expensive_check=False,
-        with_edge_properties=with_edge_properties,
-        prior_sources_behavior=prior_sources_behavior,
-        deduplicate_sources=deduplicate_sources,
-        return_hops=return_hops,
-        renumber=renumber,
-        retain_seeds=retain_seeds,
-        compression=compression,
-        compress_per_hop=compress_per_hop,
-        random_state=random_state,
-        return_dict=True,
-    )
-
-    dfs = sampling_results_from_cupy_array_dict(
-        sampling_result_array_dict,
-        weight_t,
-        len(fanout_vals),
-        with_edge_properties=with_edge_properties,
-        return_offsets=return_offsets,
-        renumber=renumber,
-        use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column,  # Deprecated flag
-    )
-
-    if G.renumbered and not renumber:
-        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
-        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
-
-    if len(dfs) > 1:
-        return dfs
-
-    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
deleted file mode 100644
index 27244ba0f96..00000000000
--- a/python/cugraph/cugraph/sampling/homogeneous_biased_uniform_neighbor_sample.py
+++ /dev/null
@@ -1,331 +0,0 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from pylibcugraph import ResourceHandle
-from pylibcugraph import (
-    homogeneous_biased_neighbor_sample as \
-        pylibcugraph_homogeneous_biased_neighbor_sample,
-)
-
-from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
-
-import numpy
-
-import cudf
-import cupy as cp
-import warnings
-
-from typing import Union, Tuple, Sequence, List
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from cugraph import Graph
-
-
-start_col_name = "_START_"
-batch_col_name = "_BATCH_"
-
-
-# FIXME: Move this function to the utility module so that it can be
-# shared by other algos
-def ensure_valid_dtype(input_graph, start_vertex_list):
-    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
-    if isinstance(start_vertex_list, cudf.Series):
-        start_vertex_list_dtypes = start_vertex_list.dtype
-    else:
-        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
-
-    if start_vertex_list_dtypes != vertex_dtype:
-        warning_msg = (
-            "biased neighbor sample requires 'start_vertex_list' to match the graph's "
-            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
-            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
-        )
-        warnings.warn(warning_msg, UserWarning)
-        start_vertex_list = start_vertex_list.astype(vertex_dtype)
-
-    return start_vertex_list
-
-
-def homogeneous_biased_neighbor_sample(
-    G: Graph,
-    start_vertex_list: Sequence,
-    start_vertex_offsets: Sequence,
-    fanout_vals: List[int],
-    with_replacement: bool = True,
-    with_edge_properties: bool = False,  # deprecated
-    prior_sources_behavior: str = None,
-    deduplicate_sources: bool = False,
-    return_hops: bool = True,
-    renumber: bool = False,
-    retain_seeds: bool = False,
-    compression: str = "COO",
-    compress_per_hop: bool = False,
-    random_state: int = None,
-    return_offsets: bool = False,
-) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
-    """
-    Does neighborhood sampling, which samples nodes from a graph based on the
-    current node's neighbors, with a corresponding fanout value at each hop.
-    The edges are sample with biases.
-
-    Parameters
-    ----------
-    G : cugraph.Graph
-        cuGraph graph, which contains connectivity information as dask cudf
-        edge list dataframe
-
-    start_vertex_list : list or cudf.Series (int32)
-        a list of starting vertices for sampling
-
-    start_vertex_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
-    fanout_vals : list (int32)
-        List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
-
-    with_replacement: bool, optional (default=True)
-        Flag to specify if the random sampling is done with replacement
-
-    with_edge_properties: bool, optional (default=False)
-        Deprecated.
-        Flag to specify whether to return edge properties (weight, edge id,
-        edge type, batch id, hop id) with the sampled edges.
-
-    prior_sources_behavior: str, optional (default=None)
-        Options are "carryover", and "exclude".
-        Default will leave the source list as-is.
-        Carryover will carry over sources from previous hops to the
-        current hop.
-        Exclude will exclude sources from previous hops from reappearing
-        as sources in future hops.
-
-    deduplicate_sources: bool, optional (default=False)
-        Whether to first deduplicate the list of possible sources
-        from the previous destinations before performing next
-        hop.
-
-    return_hops: bool, optional (default=True)
-        Whether to return the sampling results with hop ids
-        corresponding to the hop where the edge appeared.
-        Defaults to True.
-
-    renumber: bool, optional (default=False)
-        Whether to renumber on a per-batch basis.  If True,
-        will return the renumber map and renumber map offsets
-        as an additional dataframe.
-
-    retain_seeds: bool, optional (default=False)
-        If True, will retain the original seeds (original source vertices)
-        in the output even if they do not have outgoing neighbors.
-
-    compression: str, optional (default=COO)
-        Sets the compression type for the output minibatches.
-        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-
-    compress_per_hop: bool, optional (default=False)
-        Whether to compress globally (default), or to produce a separate
-        compressed edgelist per hop.
-
-    random_state: int, optional
-        Random seed to use when making sampling calls.
-
-    return_offsets: bool, optional (default=False)
-        Whether to return the sampling results with batch ids
-        included as one dataframe, or to instead return two
-        dataframes, one with sampling results and one with
-        batch ids and their start offsets.
-
-    Returns
-    -------
-    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
-        GPU data frame containing multiple cudf.Series
-
-        If with_edge_properties=False:
-            df['sources']: cudf.Series
-                Contains the source vertices from the sampling result
-            df['destinations']: cudf.Series
-                Contains the destination vertices from the sampling result
-            df['indices']: cudf.Series
-                Contains the indices (edge weights) from the sampling result
-                for path reconstruction
-
-        If with_edge_properties=True:
-            If return_offsets=False:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-
-            If return_offsets=True:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-
-                offsets_df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                offsets_df['offsets']: cudf.Series
-                    Contains the offsets of each batch in the sampling result
-
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-    """
-
-    use_legacy_names = False  # Deprecated parameter
-    include_hop_column = not return_offsets  # Deprecated parameter
-
-    major_col_name = "majors"
-    minor_col_name = "minors"
-
-    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
-        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
-
-    if (
-        (compression != "COO")
-        and (not compress_per_hop)
-        and prior_sources_behavior != "exclude"
-    ):
-        raise ValueError(
-            "hop-agnostic compression is only supported with"
-            " the exclude prior sources behavior due to limitations "
-            "of the libcugraph C++ API"
-        )
-
-    if compress_per_hop and prior_sources_behavior != "carryover":
-        raise ValueError(
-            "Compressing the edgelist per hop is only supported "
-            "with the carryover prior sources behavior due to limitations"
-            " of the libcugraph C++ API"
-        )
-
-    if with_edge_properties:
-        warning_msg = (
-            "The with_edge_properties flag is deprecated"
-            " and will be removed in the next release in favor"
-            " of returning all properties in the graph"
-        )
-        warnings.warn(warning_msg, FutureWarning)
-
-    if isinstance(start_vertex_list, int):
-        start_vertex_list = [start_vertex_list]
-
-    if isinstance(start_vertex_list, list):
-        start_vertex_list = cudf.Series(
-            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
-        )
-
-    """
-    # No batch_ids, the rank owning the vertices will wom the final
-    # result.
-    if with_edge_properties and not with_batch_ids:
-        if isinstance(start_vertex_list, cudf.Series):
-            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
-
-        start_vertex_list[batch_col_name] = cudf.Series(
-            cp.zeros(len(start_vertex_list), dtype="int32")
-        )
-    """
-
-    # fanout_vals must be passed to pylibcugraph as a host array
-    if isinstance(fanout_vals, numpy.ndarray):
-        fanout_vals = fanout_vals.astype("int32")
-    elif isinstance(fanout_vals, list):
-        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
-    elif isinstance(fanout_vals, cp.ndarray):
-        fanout_vals = fanout_vals.get().astype("int32")
-    elif isinstance(fanout_vals, cudf.Series):
-        fanout_vals = fanout_vals.values_host.astype("int32")
-    else:
-        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
-
-    if "weights" in G.edgelist.edgelist_df:
-        weight_t = G.edgelist.edgelist_df["weights"].dtype
-    else:
-        weight_t = "float32"
-
-    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
-
-    if G.renumbered:
-        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
-
-    sampling_result_array_dict = pylibcugraph_homogeneous_biased_neighbor_sample(
-        resource_handle=ResourceHandle(),
-        input_graph=G._plc_graph,
-        start_vertex_list=start_vertex_list,
-        start_vertex_offsets=start_vertex_offsets,
-        h_fan_out=fanout_vals,
-        with_replacement=with_replacement,
-        do_expensive_check=False,
-        with_edge_properties=with_edge_properties,
-        prior_sources_behavior=prior_sources_behavior,
-        deduplicate_sources=deduplicate_sources,
-        return_hops=return_hops,
-        renumber=renumber,
-        retain_seeds=retain_seeds,
-        compression=compression,
-        compress_per_hop=compress_per_hop,
-        random_state=random_state,
-        return_dict=True,
-    )
-
-    dfs = sampling_results_from_cupy_array_dict(
-        sampling_result_array_dict,
-        weight_t,
-        len(fanout_vals),
-        with_edge_properties=with_edge_properties,
-        return_offsets=return_offsets,
-        renumber=renumber,
-        use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column,  # Deprecated flag
-    )
-
-    if G.renumbered and not renumber:
-        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
-        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
-
-    if len(dfs) > 1:
-        return dfs
-
-    return dfs[0]
diff --git a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py b/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
deleted file mode 100644
index b980e38407c..00000000000
--- a/python/cugraph/cugraph/sampling/homogeneous_uniform_neighbor_sample.py
+++ /dev/null
@@ -1,330 +0,0 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from pylibcugraph import ResourceHandle
-from pylibcugraph import (
-    homogeneous_uniform_neighbor_sample as \
-        pylibcugraph_homogeneous_uniform_neighbor_sample,
-)
-
-from cugraph.sampling.sampling_utilities import sampling_results_from_cupy_array_dict
-
-import numpy
-
-import cudf
-import cupy as cp
-import warnings
-
-from typing import Union, Tuple, Sequence, List
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from cugraph import Graph
-
-
-start_col_name = "_START_"
-batch_col_name = "_BATCH_"
-
-
-# FIXME: Move this function to the utility module so that it can be
-# shared by other algos
-def ensure_valid_dtype(input_graph, start_vertex_list):
-    vertex_dtype = input_graph.edgelist.edgelist_df.dtypes.iloc[0]
-    if isinstance(start_vertex_list, cudf.Series):
-        start_vertex_list_dtypes = start_vertex_list.dtype
-    else:
-        start_vertex_list_dtypes = start_vertex_list.dtypes.iloc[0]
-
-    if start_vertex_list_dtypes != vertex_dtype:
-        warning_msg = (
-            "Uniform neighbor sample requires 'start_vertex_list' to match the graph's "
-            f"'vertex' type. input graph's vertex type is: {vertex_dtype} and got "
-            f"'start_vertex_list' of type: {start_vertex_list_dtypes}."
-        )
-        warnings.warn(warning_msg, UserWarning)
-        start_vertex_list = start_vertex_list.astype(vertex_dtype)
-
-    return start_vertex_list
-
-
-def homogeneous_uniform_neighbor_sample(
-    G: Graph,
-    start_vertex_list: Sequence,
-    start_vertex_offsets: Sequence,
-    fanout_vals: List[int],
-    with_replacement: bool = True,
-    with_edge_properties: bool = False,  # deprecated
-    prior_sources_behavior: str = None,
-    deduplicate_sources: bool = False,
-    return_hops: bool = True,
-    renumber: bool = False,
-    retain_seeds: bool = False,
-    compression: str = "COO",
-    compress_per_hop: bool = False,
-    random_state: int = None,
-    return_offsets: bool = False,
-) -> Union[cudf.DataFrame, Tuple[cudf.DataFrame, cudf.DataFrame]]:
-    """
-    Does neighborhood sampling, which samples nodes from a graph based on the
-    current node's neighbors, with a corresponding fanout value at each hop.
-
-    Parameters
-    ----------
-    G : cugraph.Graph
-        cuGraph graph, which contains connectivity information as dask cudf
-        edge list dataframe
-
-    start_vertex_list : list or cudf.Series (int32)
-        a list of starting vertices for sampling
-
-    start_vertex_offsets: list[int] (Optional)
-        Offsets of each label within the start vertex list.
-
-    fanout_vals : list (int32)
-        List of branching out (fan-out) degrees per starting vertex for each
-        hop level.
-
-    with_replacement: bool, optional (default=True)
-        Flag to specify if the random sampling is done with replacement
-
-    with_edge_properties: bool, optional (default=False)
-        Deprecated.
-        Flag to specify whether to return edge properties (weight, edge id,
-        edge type, batch id, hop id) with the sampled edges.
-
-    prior_sources_behavior: str, optional (default=None)
-        Options are "carryover", and "exclude".
-        Default will leave the source list as-is.
-        Carryover will carry over sources from previous hops to the
-        current hop.
-        Exclude will exclude sources from previous hops from reappearing
-        as sources in future hops.
-
-    deduplicate_sources: bool, optional (default=False)
-        Whether to first deduplicate the list of possible sources
-        from the previous destinations before performing next
-        hop.
-
-    return_hops: bool, optional (default=True)
-        Whether to return the sampling results with hop ids
-        corresponding to the hop where the edge appeared.
-        Defaults to True.
-
-    renumber: bool, optional (default=False)
-        Whether to renumber on a per-batch basis.  If True,
-        will return the renumber map and renumber map offsets
-        as an additional dataframe.
-
-    retain_seeds: bool, optional (default=False)
-        If True, will retain the original seeds (original source vertices)
-        in the output even if they do not have outgoing neighbors.
-
-    compression: str, optional (default=COO)
-        Sets the compression type for the output minibatches.
-        Valid options are COO (default), CSR, CSC, DCSR, and DCSC.
-
-    compress_per_hop: bool, optional (default=False)
-        Whether to compress globally (default), or to produce a separate
-        compressed edgelist per hop.
-
-    random_state: int, optional
-        Random seed to use when making sampling calls.
-
-    return_offsets: bool, optional (default=False)
-        Whether to return the sampling results with batch ids
-        included as one dataframe, or to instead return two
-        dataframes, one with sampling results and one with
-        batch ids and their start offsets.
-
-    Returns
-    -------
-    result : cudf.DataFrame or Tuple[cudf.DataFrame, cudf.DataFrame]
-        GPU data frame containing multiple cudf.Series
-
-        If with_edge_properties=False:
-            df['sources']: cudf.Series
-                Contains the source vertices from the sampling result
-            df['destinations']: cudf.Series
-                Contains the destination vertices from the sampling result
-            df['indices']: cudf.Series
-                Contains the indices (edge weights) from the sampling result
-                for path reconstruction
-
-        If with_edge_properties=True:
-            If return_offsets=False:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-
-            If return_offsets=True:
-                df['sources']: cudf.Series
-                    Contains the source vertices from the sampling result
-                df['destinations']: cudf.Series
-                    Contains the destination vertices from the sampling result
-                df['edge_weight']: cudf.Series
-                    Contains the edge weights from the sampling result
-                df['edge_id']: cudf.Series
-                    Contains the edge ids from the sampling result
-                df['edge_type']: cudf.Series
-                    Contains the edge types from the sampling result
-                df['hop_id']: cudf.Series
-                    Contains the hop ids from the sampling result
-
-                offsets_df['batch_id']: cudf.Series
-                    Contains the batch ids from the sampling result
-                offsets_df['offsets']: cudf.Series
-                    Contains the offsets of each batch in the sampling result
-
-                If renumber=True:
-                    (adds the following dataframe)
-                    renumber_df['map']: cudf.Series
-                        Contains the renumber maps for each batch
-                    renumber_df['offsets']: cudf.Series
-                        Contains the batch offsets for the renumber maps
-    """
-
-    use_legacy_names = False  # Deprecated parameter
-    include_hop_column = not return_offsets  # Deprecated parameter
-
-    major_col_name = "majors"
-    minor_col_name = "minors"
-
-    if compression not in ["COO", "CSR", "CSC", "DCSR", "DCSC"]:
-        raise ValueError("compression must be one of COO, CSR, CSC, DCSR, or DCSC")
-
-    if (
-        (compression != "COO")
-        and (not compress_per_hop)
-        and prior_sources_behavior != "exclude"
-    ):
-        raise ValueError(
-            "hop-agnostic compression is only supported with"
-            " the exclude prior sources behavior due to limitations "
-            "of the libcugraph C++ API"
-        )
-
-    if compress_per_hop and prior_sources_behavior != "carryover":
-        raise ValueError(
-            "Compressing the edgelist per hop is only supported "
-            "with the carryover prior sources behavior due to limitations"
-            " of the libcugraph C++ API"
-        )
-
-    if with_edge_properties:
-        warning_msg = (
-            "The with_edge_properties flag is deprecated"
-            " and will be removed in the next release in favor"
-            " of returning all properties in the graph"
-        )
-        warnings.warn(warning_msg, FutureWarning)
-
-    if isinstance(start_vertex_list, int):
-        start_vertex_list = [start_vertex_list]
-
-    if isinstance(start_vertex_list, list):
-        start_vertex_list = cudf.Series(
-            start_vertex_list, dtype=G.edgelist.edgelist_df[G.srcCol].dtype
-        )
-
-    """
-    # No batch_ids, the rank owning the vertices will wom the final
-    # result.
-    if with_edge_properties and not with_batch_ids:
-        if isinstance(start_vertex_list, cudf.Series):
-            start_vertex_list = start_vertex_list.reset_index(drop=True).to_frame()
-
-        start_vertex_list[batch_col_name] = cudf.Series(
-            cp.zeros(len(start_vertex_list), dtype="int32")
-        )
-    """
-
-    # fanout_vals must be passed to pylibcugraph as a host array
-    if isinstance(fanout_vals, numpy.ndarray):
-        fanout_vals = fanout_vals.astype("int32")
-    elif isinstance(fanout_vals, list):
-        fanout_vals = numpy.asarray(fanout_vals, dtype="int32")
-    elif isinstance(fanout_vals, cp.ndarray):
-        fanout_vals = fanout_vals.get().astype("int32")
-    elif isinstance(fanout_vals, cudf.Series):
-        fanout_vals = fanout_vals.values_host.astype("int32")
-    else:
-        raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
-
-    if "weights" in G.edgelist.edgelist_df:
-        weight_t = G.edgelist.edgelist_df["weights"].dtype
-    else:
-        weight_t = "float32"
-
-    start_vertex_list = ensure_valid_dtype(G, start_vertex_list)
-
-    if G.renumbered:
-        start_vertex_list = G.lookup_internal_vertex_id(start_vertex_list)
-
-    sampling_result_array_dict = pylibcugraph_homogeneous_uniform_neighbor_sample(
-        resource_handle=ResourceHandle(),
-        input_graph=G._plc_graph,
-        start_vertex_list=start_vertex_list,
-        start_vertex_offsets=start_vertex_offsets,
-        h_fan_out=fanout_vals,
-        with_replacement=with_replacement,
-        do_expensive_check=False,
-        with_edge_properties=with_edge_properties,
-        prior_sources_behavior=prior_sources_behavior,
-        deduplicate_sources=deduplicate_sources,
-        return_hops=return_hops,
-        renumber=renumber,
-        retain_seeds=retain_seeds,
-        compression=compression,
-        compress_per_hop=compress_per_hop,
-        random_state=random_state,
-        return_dict=True,
-    )
-
-    dfs = sampling_results_from_cupy_array_dict(
-        sampling_result_array_dict,
-        weight_t,
-        len(fanout_vals),
-        with_edge_properties=with_edge_properties,
-        return_offsets=return_offsets,
-        renumber=renumber,
-        use_legacy_names=use_legacy_names,
-        include_hop_column=include_hop_column,  # Deprecated flag
-    )
-
-    if G.renumbered and not renumber:
-        dfs[0] = G.unrenumber(dfs[0], major_col_name, preserve_order=True)
-        dfs[0] = G.unrenumber(dfs[0], minor_col_name, preserve_order=True)
-
-    if len(dfs) > 1:
-        return dfs
-
-    return dfs[0]

From 5b1bbb49cefe18a7d791eb4a84dbcb5855c7729c Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 09:50:48 -0700
Subject: [PATCH 090/117] undo changes to uniform neighbor sampling

---
 .../dask/sampling/uniform_neighbor_sample.py  | 27 +++----------------
 1 file changed, 3 insertions(+), 24 deletions(-)

diff --git a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
index 36a86635254..03746561817 100644
--- a/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
+++ b/python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2024, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,9 +13,6 @@
 # limitations under the License.
 
 from __future__ import annotations
-from itertools import chain
-import collections
-from itertools import accumulate
 
 import warnings
 
@@ -380,10 +377,9 @@ def uniform_neighbor_sample(
     start_list : int, list, cudf.Series, or dask_cudf.Series (int32 or int64)
         a list of starting vertices for sampling
 
-    fanout_vals : list (int32) or dict
+    fanout_vals : list
         List of branching out (fan-out) degrees per starting vertex for each
-        hop level or dictionary of edge type and fanout values for
-        heterogeneous fanout type.
+        hop level.
 
     with_replacement: bool, optional (default=True)
         Flag to specify if the random sampling is done with replacement
@@ -616,23 +612,6 @@ def uniform_neighbor_sample(
         fanout_vals = fanout_vals.get().astype("int32")
     elif isinstance(fanout_vals, cudf.Series):
         fanout_vals = fanout_vals.values_host.astype("int32")
-    elif isinstance(fanout_vals, dict):
-        # FIXME: Add expensive check to ensure all dict values are lists
-        # Convert to a tuple of sequence (edge type size and fanout values)
-
-        fanout_vals = {2: [1], 3: [4, 5], 1: [6, 7], 0: [8, 9, 10]}
-        fanout_vals = collections.OrderedDict(sorted(fanout_vals.items()))
-
-        edge_type_offsets = []
-        [edge_type_offsets.append(len(s)) for s in list(fanout_vals.values())]
-
-        edge_type_offsets = list(accumulate(edge_type_offsets))
-        edge_type_fanout_vals = list(chain.from_iterable(list(fanout_vals.values())))
-
-        fanout_vals = (
-            numpy.asarray(edge_type_offsets, dtype="int32"),
-            numpy.asarray(edge_type_fanout_vals, dtype="int32"),
-        )
     else:
         raise TypeError("fanout_vals must be a sequence, " f"got: {type(fanout_vals)}")
 

From 6d69f88dcab0d7b2da32d1a6674e6d4f80bd6299 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 09:52:22 -0700
Subject: [PATCH 091/117] remove o utdated fixme

---
 python/pylibcugraph/pylibcugraph/__init__.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py
index 97c8ab342e6..b9a981094c8 100644
--- a/python/pylibcugraph/pylibcugraph/__init__.py
+++ b/python/pylibcugraph/pylibcugraph/__init__.py
@@ -53,11 +53,6 @@
     heterogeneous_biased_neighbor_sample,
 )
 
-
-# FIXME: break down the API into homogeneous and heterogeneous neighbor sample
-# from pylibcugraph.uniform_neighbor_sample import neighbor_sample
-# from pylibcugraph.biased_neighbor_sample import biased_neighbor_sample
-
 from pylibcugraph.negative_sampling import negative_sampling
 
 from pylibcugraph.core_number import core_number

From 79d45270e50172d57a14cbb39efb5a316b834927 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 09:58:48 -0700
Subject: [PATCH 092/117] remove unnecessary call

---
 ...heterogeneous_biased_neighbor_sampling.cpp | 30 +------------------
 ...eterogeneous_uniform_neighbor_sampling.cpp | 29 ------------------
 .../homogeneous_biased_neighbor_sampling.cpp  | 24 ---------------
 .../homogeneous_uniform_neighbor_sampling.cpp | 23 --------------
 ...heterogeneous_biased_neighbor_sampling.cpp | 28 -----------------
 ...eterogeneous_uniform_neighbor_sampling.cpp | 29 ------------------
 ...g_homogeneous_biased_neighbor_sampling.cpp | 24 ---------------
 ..._homogeneous_uniform_neighbor_sampling.cpp | 23 --------------
 8 files changed, 1 insertion(+), 209 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index 3c3739ed668..995728755a3 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -124,34 +124,7 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         handle, graph_view, heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
     }
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::heterogeneous_biased_neighbor_sample(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        edge_types
-          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                       .view()}
-          : std::nullopt,
-        *edge_weight_view,
-        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
-        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
-                                                                           batch_number->size()})
-                     : std::nullopt,
-        label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(
-          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  heterogeneous_biased_neighbor_sampling_usecase.flag_replacement}),
-      std::exception);
-#else
+
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.start("Biased neighbor sampling");
@@ -233,7 +206,6 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
           heterogeneous_biased_neighbor_sampling_usecase.fanout.size()));
       }
     }
-#endif
   }
 };
 
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 83f99e68e9a..b09ff329e50 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -124,34 +124,6 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
         handle, graph_view, heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
     }
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::heterogeneous_uniform_neighbor_sample(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        edge_types
-          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                       .view()}
-          : std::nullopt,
-        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
-        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
-                                                                           batch_number->size()})
-                     : std::nullopt,
-        label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(
-          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-          cugraph::prior_sources_behavior_t{0},
-          true,   // return_hops
-          false,  // dedupe_sources
-          heterogeneous_uniform_neighbor_sampling_usecase.flag_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.start("Uniform neighbor sampling");
@@ -233,7 +205,6 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
           heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
       }
     }
-#endif
   }
 };
 
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
index d01206c5cef..3e4d8ea2adc 100644
--- a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -118,29 +118,6 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
 
     std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::homogeneous_biased_neighbor_sample(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        *edge_weight_view,
-        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
-        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
-                                                                           batch_number->size()})
-                     : std::nullopt,
-        label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  homogeneous_biased_neighbor_sampling_usecase.flag_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.start("Biased neighbor sampling");
@@ -217,7 +194,6 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
           homogeneous_biased_neighbor_sampling_usecase.fanout.size()));
       }
     }
-#endif
   }
 };
 
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
index 508914e2495..337927e5580 100644
--- a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -123,28 +123,6 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
 
     std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank_mapping{std::nullopt};
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::homogeneous_uniform_neighbor_sample(
-        handle,
-        rng_state,
-        graph_view,
-        edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        raft::device_span<vertex_t const>{random_sources_copy.data(), random_sources.size()},
-        batch_number ? std::make_optional(raft::device_span<int32_t const>{batch_number->data(),
-                                                                           batch_number->size()})
-                     : std::nullopt,
-        label_to_output_comm_rank_mapping,
-        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  homogeneous_uniform_neighbor_sampling_usecase.flag_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.start("Uniform neighbor sampling");
@@ -220,7 +198,6 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
           homogeneous_uniform_neighbor_sampling_usecase.fanout.size()));
       }
     }
-#endif
   }
 };
 
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index 688d8480118..35461f6ed7d 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -159,33 +159,6 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
         *handle_, mg_graph_view, heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
     }
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      ccugraph::heterogeneous_biased_neighbor_sample(
-        *handle_,
-        rng_state,
-        mg_graph_view,
-        mg_edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        edge_types
-          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                       .view()}
-          : std::nullopt,
-
-        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
-        std::make_optional(
-          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(
-          heterogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-          heterogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-        heterogeneous_biased_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  heterogeneous_biased_neighbor_sampling_usecase.with_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       handle_->get_comms().barrier();
@@ -309,7 +282,6 @@ class Tests_MGHeterogeneous_Biased_Neighbor_Sampling
         }
       }
     }
-#endif
   }
 
  private:
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index 18e75d1b1b6..ed085c67f6a 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -159,34 +159,6 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
         *handle_, mg_graph_view, heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types);
     }
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      ccugraph::heterogeneous_uniform_neighbor_sample(
-        *handle_,
-        rng_state,
-        mg_graph_view,
-        mg_edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        edge_types
-          ? std::optional<cugraph::edge_property_view_t<edge_t, edge_type_t const*>>{(*edge_types)
-                                                                                       .view()}
-          : std::nullopt,
-
-        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
-        std::make_optional(
-          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(
-          heterogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-          heterogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-        heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types,
-        cugraph::sampling_flags_t{
-          cugraph::prior_sources_behavior_t{0},
-          true,   // return_hops
-          false,  // dedupe_sources
-          heterogeneous_uniform_neighbor_sampling_usecase.with_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       handle_->get_comms().barrier();
@@ -310,7 +282,6 @@ class Tests_MGHeterogeneous_Uniform_Neighbor_Sampling
         }
       }
     }
-#endif
   }
 
  private:
diff --git a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
index ec134938067..701816645f9 100644
--- a/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_biased_neighbor_sampling.cpp
@@ -153,29 +153,6 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
                random_sources.size(),
                handle_->get_stream());
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::homogeneous_biased_neighbor_sample(
-        *handle_,
-        rng_state,
-        mg_graph_view,
-        mg_edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        mg_edge_weight_view raft::device_span<vertex_t const>{random_sources.data(),
-                                                              random_sources.size()},
-        std::make_optional(
-          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(homogeneous_biased_neighbor_sampling_usecase.fanout.data(),
-                                       homogeneous_biased_neighbor_sampling_usecase.fanout.size()),
-
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  homogeneous_biased_neighbor_sampling_usecase.with_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       handle_->get_comms().barrier();
@@ -294,7 +271,6 @@ class Tests_MGHomogeneous_Biased_Neighbor_Sampling
         }
       }
     }
-#endif
   }
 
  private:
diff --git a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
index e33000044a5..4d9d82fc719 100644
--- a/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_homogeneous_uniform_neighbor_sampling.cpp
@@ -154,28 +154,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
                random_sources.size(),
                handle_->get_stream());
 
-#ifdef NO_CUGRAPH_OPS
-    EXPECT_THROW(
-      cugraph::homogeneous_uniform_neighbor_sample(
-        *handle_,
-        rng_state,
-        mg_graph_view,
-        mg_edge_weight_view,
-        std::optional<cugraph::edge_property_view_t<edge_t, edge_t const*>>{std::nullopt},
-        std::optional<cugraph::edge_property_view_t<edge_t, int32_t const*>>{std::nullopt},
-        raft::device_span<vertex_t const>{random_sources.data(), random_sources.size()},
-        std::make_optional(
-          raft::device_span<int32_t const>{batch_number.data(), batch_number.size()}),
-        std::make_optional(raft::device_span<int32_t const>{comm_ranks.data(), comm_ranks.size()}),
-        raft::host_span<int32_t const>(homogeneous_uniform_neighbor_sampling_usecase.fanout.data(),
-                                       homogeneous_uniform_neighbor_sampling_usecase.fanout.size()),
-
-        cugraph::sampling_flags_t{cugraph::prior_sources_behavior_t{0},
-                                  true,   // return_hops
-                                  false,  // dedupe_sources
-                                  homogeneous_uniform_neighbor_sampling_usecase.with_replacement}),
-      std::exception);
-#else
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       handle_->get_comms().barrier();
@@ -294,7 +272,6 @@ class Tests_MGHomogeneous_Uniform_Neighbor_Sampling
         }
       }
     }
-#endif
   }
 
  private:

From 2a6692818c9432c02e6e99ded0139938c29998f3 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:01:47 -0700
Subject: [PATCH 093/117] remove unnecessary blank line

---
 cpp/include/cugraph_c/sampling_algorithms.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index fb0b93a0bed..dde2e270441 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -863,7 +863,6 @@ cugraph_error_code_t cugraph_test_uniform_neighborhood_sample_result_create(
  *                                not CUGRAPH_SUCCESS
  * @return error code
  */
-
 cugraph_error_code_t cugraph_select_random_vertices(const cugraph_resource_handle_t* handle,
                                                     const cugraph_graph_t* graph,
                                                     cugraph_rng_state_t* rng_state,

From 8c3d87150edbf9c42825a7a92b0113bda23ce406 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:05:19 -0700
Subject: [PATCH 094/117] add comments for deprecated functions

---
 cpp/src/c_api/neighbor_sampling.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 527fa53ff65..f4384ff7bbc 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -1880,6 +1880,7 @@ cugraph_error_code_t cugraph_uniform_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
+  //  Deprecated functor
   uniform_neighbor_sampling_functor functor{handle,
                                             graph,
                                             start_vertices,
@@ -1957,6 +1958,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
     "fan_out should be of type int",
     *error);
 
+  // Deprecated functor
   biased_neighbor_sampling_functor functor{handle,
                                            graph,
                                            edge_biases,

From ae92c9f51c6907da3fc13415b84a73ff9cf3538b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:08:38 -0700
Subject: [PATCH 095/117] remove obsolete instantiation

---
 .../shuffle_and_organize_output_mg_v32_e64.cu |  61 ----
 .../sampling/neighbor_sampling_mg_v32_e64.cu  | 307 ------------------
 2 files changed, 368 deletions(-)
 delete mode 100644 cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
 delete mode 100644 cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu

diff --git a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu b/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
deleted file mode 100644
index ef760844b51..00000000000
--- a/cpp/src/sampling/detail/shuffle_and_organize_output_mg_v32_e64.cu
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "sampling/detail/shuffle_and_organize_output_impl.cuh"
-
-namespace cugraph {
-namespace detail {
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-shuffle_and_organize_output(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& majors,
-  rmm::device_uvector<int32_t>&& minors,
-  std::optional<rmm::device_uvector<float>>&& weights,
-  std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
-  std::optional<rmm::device_uvector<int32_t>>&& edge_types,
-  std::optional<rmm::device_uvector<int32_t>>&& hops,
-  std::optional<rmm::device_uvector<int32_t>>&& labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-shuffle_and_organize_output(
-  raft::handle_t const& handle,
-  rmm::device_uvector<int32_t>&& majors,
-  rmm::device_uvector<int32_t>&& minors,
-  std::optional<rmm::device_uvector<double>>&& weights,
-  std::optional<rmm::device_uvector<int64_t>>&& edge_ids,
-  std::optional<rmm::device_uvector<int32_t>>&& edge_types,
-  std::optional<rmm::device_uvector<int32_t>>&& hops,
-  std::optional<rmm::device_uvector<int32_t>>&& labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank);
-
-}  // namespace detail
-}  // namespace cugraph
diff --git a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
deleted file mode 100644
index 2a387405900..00000000000
--- a/cpp/src/sampling/neighbor_sampling_mg_v32_e64.cu
+++ /dev/null
@@ -1,307 +0,0 @@
-/*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "neighbor_sampling_impl.hpp"
-
-#include <cugraph/algorithms.hpp>
-#include <cugraph/sampling_functions.hpp>
-
-namespace cugraph {
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-biased_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-biased_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, true> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-}  // namespace cugraph

From d7d6109d3c6a4da6633e63e1e04d950a9c94f201 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:11:28 -0700
Subject: [PATCH 096/117] remove unnecessary parenthesis

---
 cpp/src/sampling/neighbor_sampling_impl.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/sampling/neighbor_sampling_impl.hpp b/cpp/src/sampling/neighbor_sampling_impl.hpp
index e3a06f323d1..935bdfceca3 100644
--- a/cpp/src/sampling/neighbor_sampling_impl.hpp
+++ b/cpp/src/sampling/neighbor_sampling_impl.hpp
@@ -339,7 +339,7 @@ neighbor_sample_impl(raft::handle_t const& handle,
   if (return_hops) {
     result_hops   = rmm::device_uvector<int32_t>(result_size, handle.get_stream());
     output_offset = 0;
-    for (size_t i = 0; i < (fan_out).size(); ++i) {
+    for (size_t i = 0; i < fan_out.size(); ++i) {
       scalar_fill(
         handle, result_hops->data() + output_offset, level_sizes[i], static_cast<int32_t>(i));
       output_offset += level_sizes[i];

From 6b3ffbd34703714a32534d65cf2b0c9df59f9b8e Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:13:11 -0700
Subject: [PATCH 097/117] remove obsolete instantiation

---
 .../sampling/neighbor_sampling_sg_v32_e64.cu  | 306 ------------------
 1 file changed, 306 deletions(-)
 delete mode 100644 cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu

diff --git a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu b/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
deleted file mode 100644
index 47cb064cd37..00000000000
--- a/cpp/src/sampling/neighbor_sampling_sg_v32_e64.cu
+++ /dev/null
@@ -1,306 +0,0 @@
-/*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "neighbor_sampling_impl.hpp"
-
-#include <cugraph/algorithms.hpp>
-#include <cugraph/sampling_functions.hpp>
-
-namespace cugraph {
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-biased_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-biased_neighbor_sample(
-  raft::handle_t const& handle,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  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,
-  raft::host_span<int32_t const> fan_out,
-  raft::random::RngState& rng_state,
-  bool return_hops,
-  bool with_replacement,
-  prior_sources_behavior_t prior_sources_behavior,
-  bool dedupe_sources,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-heterogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  int32_t num_edge_types,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_uniform_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<double>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, double const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, double const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-template std::tuple<rmm::device_uvector<int32_t>,
-                    rmm::device_uvector<int32_t>,
-                    std::optional<rmm::device_uvector<float>>,
-                    std::optional<rmm::device_uvector<int64_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<int32_t>>,
-                    std::optional<rmm::device_uvector<size_t>>>
-homogeneous_biased_neighbor_sample(
-  raft::handle_t const& handle,
-  raft::random::RngState& rng_state,
-  graph_view_t<int32_t, int64_t, false, false> const& graph_view,
-  std::optional<edge_property_view_t<int64_t, float const*>> edge_weight_view,
-  std::optional<edge_property_view_t<int64_t, int64_t const*>> edge_id_view,
-  std::optional<edge_property_view_t<int64_t, int32_t const*>> edge_type_view,
-  edge_property_view_t<int64_t, float const*> edge_bias_view,
-  raft::device_span<int32_t const> starting_vertices,
-  std::optional<raft::device_span<int32_t const>> starting_vertex_labels,
-  std::optional<raft::device_span<int32_t const>> label_to_output_comm_rank,
-  raft::host_span<int32_t const> fan_out,
-  sampling_flags_t sampling_flags,
-  bool do_expensive_check);
-
-}  // namespace cugraph

From 67d7d0a1a623017b87a53248fb4da75badd463ac Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 10:15:49 -0700
Subject: [PATCH 098/117] fix style

---
 .../cugraph/detail/utility_wrappers.hpp       |  4 +-
 cpp/src/c_api/neighbor_sampling.cpp           | 17 +++++----
 cpp/src/detail/utility_wrappers_impl.cuh      | 37 +++++++++----------
 ...heterogeneous_biased_neighbor_sampling.cpp |  1 -
 ...eterogeneous_uniform_neighbor_sampling.cpp |  3 +-
 .../homogeneous_uniform_neighbor_sampling.cpp |  7 ++--
 6 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index a9486621798..585ee4658f7 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -104,7 +104,9 @@ size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value);
  *
  */
 template <typename value_t>
-void transform_increment(rmm::cuda_stream_view const& stream_view, raft::device_span<value_t> d_value, value_t value);
+void transform_increment(rmm::cuda_stream_view const& stream_view,
+                         raft::device_span<value_t> d_value,
+                         value_t value);
 
 /**
  * @brief    Fill a buffer with a sequence of values
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index f4384ff7bbc..54a2632b7c4 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -885,10 +885,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
         // Compute the global start_vertex_label_offsets
 
-        cugraph::detail::transform_increment(handle_.get_stream(),
-                                              raft::device_span<label_t>{(*start_vertex_labels).data(), (*start_vertex_labels).size()},
-                                              (label_t)global_labels[handle_.get_comms().get_rank()]
-                                              );
+        cugraph::detail::transform_increment(
+          handle_.get_stream(),
+          raft::device_span<label_t>{(*start_vertex_labels).data(), (*start_vertex_labels).size()},
+          (label_t)global_labels[handle_.get_comms().get_rank()]);
       }
 
       if constexpr (multi_gpu) {
@@ -902,10 +902,13 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
           // Get unique labels
           // sort the start_vertex_labels
-          cugraph::detail::sort(handle_.get_stream(), raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
-          
+          cugraph::detail::sort(
+            handle_.get_stream(),
+            raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
+
           auto num_unique_labels = cugraph::detail::unique(
-            handle_.get_stream(), raft::device_span<label_t const>{unique_labels.data(), unique_labels.size()});
+            handle_.get_stream(),
+            raft::device_span<label_t const>{unique_labels.data(), unique_labels.size()});
 
           (*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index 60bb1557019..320ab65d37f 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -73,7 +73,8 @@ void sort(raft::handle_t const& handle, raft::device_span<value_t> d_value)
 template <typename value_t>
 size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value)
 {
-  auto unique_element_last = thrust::unique(handle.get_thrust_policy(), d_value.begin(), d_value.end());
+  auto unique_element_last =
+    thrust::unique(handle.get_thrust_policy(), d_value.begin(), d_value.end());
   return thrust::distance(d_value.begin(), unique_element_last);
 }
 
@@ -86,28 +87,24 @@ void sequence_fill(rmm::cuda_stream_view const& stream_view,
   thrust::sequence(rmm::exec_policy(stream_view), d_value, d_value + size, start_value);
 }
 
-
 template <typename value_t>
- void transform_increment(rmm::cuda_stream_view const& stream_view,
-                          raft::device_span<value_t> d_value,
-                          value_t incr)
- {
-   thrust::transform(rmm::exec_policy(stream_view),
-                     d_value.begin(),
-                     d_value.end(),
-                     d_value.begin(),
-                     cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
-                       return static_cast<value_t>(value + incr);
-                     }));
- }
-
+void transform_increment(rmm::cuda_stream_view const& stream_view,
+                         raft::device_span<value_t> d_value,
+                         value_t incr)
+{
+  thrust::transform(rmm::exec_policy(stream_view),
+                    d_value.begin(),
+                    d_value.end(),
+                    d_value.begin(),
+                    cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
+                      return static_cast<value_t>(value + incr);
+                    }));
+}
 
 template <typename value_t>
- void transform_increment_(rmm::cuda_stream_view const& stream_view,
-                           value_t d_value)
- {
-
- }
+void transform_increment_(rmm::cuda_stream_view const& stream_view, value_t d_value)
+{
+}
 
 template <typename value_t>
 void stride_fill(rmm::cuda_stream_view const& stream_view,
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index 995728755a3..eabcee044e0 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -124,7 +124,6 @@ class Tests_Heterogeneous_Biased_Neighbor_Sampling
         handle, graph_view, heterogeneous_biased_neighbor_sampling_usecase.num_edge_types);
     }
 
-
     if (cugraph::test::g_perf) {
       RAFT_CUDA_TRY(cudaDeviceSynchronize());  // for consistent performance measurement
       hr_timer.start("Biased neighbor sampling");
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index b09ff329e50..3f255d6b8f3 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -116,8 +116,7 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
     // Generate the edge types
 
-    std::optional<cugraph::edge_property_t<decltype(graph_view), int32_t>> edge_types{
-      std::nullopt};
+    std::optional<cugraph::edge_property_t<decltype(graph_view), int32_t>> edge_types{std::nullopt};
 
     if (heterogeneous_uniform_neighbor_sampling_usecase.num_edge_types > 1) {
       edge_types = cugraph::test::generate<decltype(graph_view), int32_t>::edge_property(
diff --git a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
index 337927e5580..a257e424b3e 100644
--- a/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_uniform_neighbor_sampling.cpp
@@ -80,10 +80,9 @@ class Tests_Homogeneous_Uniform_Neighbor_Sampling
       graph_view.attach_edge_mask((*edge_mask).view());
     }
 
-    // FIXME: Read a tuple of two edge mask and mask out if edge mask is set in either 1 (OR) and create
-    // a new one.
-    // No graph view can have two mask and perform OR in itself, and need to OR the mask 
-    // manually by itself.
+    // FIXME: Read a tuple of two edge mask and mask out if edge mask is set in either 1 (OR) and
+    // create a new one. No graph view can have two mask and perform OR in itself, and need to OR
+    // the mask manually by itself.
 
     constexpr float select_probability{0.05};
 

From 413a577bed5ea72559dc5a535d4ee4a1c784b6e9 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 11:58:12 -0700
Subject: [PATCH 099/117] fix type error

---
 cpp/src/c_api/neighbor_sampling.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 54a2632b7c4..2988c992fc4 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -908,7 +908,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
           auto num_unique_labels = cugraph::detail::unique(
             handle_.get_stream(),
-            raft::device_span<label_t const>{unique_labels.data(), unique_labels.size()});
+            raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
 
           (*local_label_to_comm_rank).resize(num_unique_labels, handle_.get_stream());
 

From 22db98d6f40a485b314193cb55ff0aff25c92dc8 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 30 Oct 2024 14:45:10 -0700
Subject: [PATCH 100/117] fix import error

---
 .../pylibcugraph/heterogeneous_biased_neighbor_sample.pyx    | 5 +++--
 .../pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx   | 5 +++--
 .../pylibcugraph/homogeneous_biased_neighbor_sample.pyx      | 5 +++--
 .../pylibcugraph/homogeneous_uniform_neighbor_sample.pyx     | 5 +++--
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
index ea71caaa7db..ba39a85f27b 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_biased_neighbor_sample.pyx
@@ -15,9 +15,10 @@
 # cython: language_level = 3
 
 from libc.stdint cimport uintptr_t
-
-from pylibcugraph._cugraph_c.resource_handle cimport (
+from pylibcugraph._cugraph_c.types cimport (
     bool_t,
+)
+from pylibcugraph._cugraph_c.resource_handle cimport (
     cugraph_resource_handle_t,
 )
 from pylibcugraph._cugraph_c.properties cimport (
diff --git a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
index 2580c7369be..8d72883e73a 100644
--- a/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/heterogeneous_uniform_neighbor_sample.pyx
@@ -15,9 +15,10 @@
 # cython: language_level = 3
 
 from libc.stdint cimport uintptr_t
-
-from pylibcugraph._cugraph_c.resource_handle cimport (
+from pylibcugraph._cugraph_c.types cimport (
     bool_t,
+)
+from pylibcugraph._cugraph_c.resource_handle cimport (
     cugraph_resource_handle_t,
 )
 from pylibcugraph._cugraph_c.error cimport (
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
index 51b770f0fab..ae5a86ce0fa 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_biased_neighbor_sample.pyx
@@ -15,9 +15,10 @@
 # cython: language_level = 3
 
 from libc.stdint cimport uintptr_t
-
-from pylibcugraph._cugraph_c.resource_handle cimport (
+from pylibcugraph._cugraph_c.types cimport (
     bool_t,
+)
+from pylibcugraph._cugraph_c.resource_handle cimport (
     cugraph_resource_handle_t,
 )
 from pylibcugraph._cugraph_c.properties cimport (
diff --git a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
index bddcc0b29ff..7c1a49df31d 100644
--- a/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
+++ b/python/pylibcugraph/pylibcugraph/homogeneous_uniform_neighbor_sample.pyx
@@ -15,9 +15,10 @@
 # cython: language_level = 3
 
 from libc.stdint cimport uintptr_t
-
-from pylibcugraph._cugraph_c.resource_handle cimport (
+from pylibcugraph._cugraph_c.types cimport (
     bool_t,
+)
+from pylibcugraph._cugraph_c.resource_handle cimport (
     cugraph_resource_handle_t,
 )
 from pylibcugraph._cugraph_c.error cimport (

From 6a95852b192655b4782fcdd015cd70eb6fd3c179 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 31 Oct 2024 09:37:43 -0700
Subject: [PATCH 101/117] remove redundant tests

---
 cpp/tests/CMakeLists.txt                         |  3 ++-
 .../heterogeneous_biased_neighbor_sampling.cpp   | 16 +++-------------
 .../heterogeneous_uniform_neighbor_sampling.cpp  | 10 ----------
 ...mg_heterogeneous_biased_neighbor_sampling.cpp | 10 +---------
 ...g_heterogeneous_uniform_neighbor_sampling.cpp | 11 +----------
 5 files changed, 7 insertions(+), 43 deletions(-)

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index f392c7fb461..a2eeafea8cf 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -503,7 +503,8 @@ ConfigureTest(
 ###################################################################################################
 # - HETEROGENEOUS BIASED NBR SAMPLING tests -------------------------------------------------------
 ConfigureTest(
-    HETEROGENEOUS_BIASED_NEIGHBOR_SAMPLING_TEST sampling/heterogeneous_biased_neighbor_sampling.cpp)
+    HETEROGENEOUS_BIASED_NEIGHBOR_SAMPLING_TESTT sampling/heterogeneous_biased_neighbor_sampling.cpp
+        GPUS 1 PERCENT 75)
 
 ###################################################################################################
 # - SAMPLING_POST_PROCESSING tests ----------------------------------------------------------------
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index eabcee044e0..1da5dc539c2 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -243,8 +243,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
@@ -253,20 +251,16 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_Heterogeneous_Biased_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
-    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
-                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
-                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+    ::testing::Values(cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/webbase-1M.mtx"))));
 
 INSTANTIATE_TEST_SUITE_P(
   rmat_small_test,
   Tests_Heterogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
 
@@ -281,10 +275,6 @@ INSTANTIATE_TEST_SUITE_P(
     ::testing::Values(
       Heterogeneous_Biased_Neighbor_Sampling_Usecase{
         {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
-        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-      Heterogeneous_Biased_Neighbor_Sampling_Usecase{
-        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
       Heterogeneous_Biased_Neighbor_Sampling_Usecase{
         {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 3f255d6b8f3..db2e3fbcb05 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -242,8 +242,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
@@ -252,8 +250,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_Heterogeneous_Uniform_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
                       cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
@@ -264,8 +260,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_Heterogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false, 0))));
 
@@ -280,10 +274,6 @@ INSTANTIATE_TEST_SUITE_P(
     ::testing::Values(
       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
         {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
-        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, false, false},
-      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
-        {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false},
       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
         {4, 10, 7, 8, 1, 9, 5, 12}, 1024, 4, true, false}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false, 0))));
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index 35461f6ed7d..e08dcb087cc 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -321,8 +321,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
@@ -331,8 +329,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Biased_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
                       cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
@@ -343,8 +339,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(
       // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
@@ -359,9 +353,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Biased_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false}),
+                      Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index ed085c67f6a..bf338b964f0 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -321,8 +321,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/karate.mtx"))));
 
@@ -331,8 +329,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_File,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
                       cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
@@ -343,8 +339,6 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(
       // cugraph::test::Rmat_Usecase(10, 16, 0.57, 0.19, 0.19, 0, false, false))));
@@ -359,10 +353,7 @@ INSTANTIATE_TEST_SUITE_P(
   Tests_MGHeterogeneous_Uniform_Neighbor_Sampling_Rmat,
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
-                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{
-                        {4, 10, 7, 8}, 128, 2, false}),
+                      Heterogeneous_Uniform_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
     ::testing::Values(cugraph::test::Rmat_Usecase(20, 32, 0.57, 0.19, 0.19, 0, false, false))));
 
 CUGRAPH_MG_TEST_PROGRAM_MAIN()

From 4f5dc3e0f1f737dcf142db3de2f7e7a10af42228 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 31 Oct 2024 09:42:42 -0700
Subject: [PATCH 102/117] remove hardcoded path

---
 .../sampling/heterogeneous_biased_neighbor_sampling.cpp    | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index 1da5dc539c2..ce0b3471cd6 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -252,9 +252,10 @@ INSTANTIATE_TEST_SUITE_P(
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
-    ::testing::Values(cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/web-Google.mtx"),
-                      cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/ljournal-2008.mtx"),
-                      cugraph::test::File_Usecase("/raid/seunghwak/datasets/test/datasets/webbase-1M.mtx"))));
+    ::testing::Values(
+      cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
 
 INSTANTIATE_TEST_SUITE_P(
   rmat_small_test,

From 9b4683a93dda2cd892c880f406a155a4cc114f9e Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 31 Oct 2024 10:04:08 -0700
Subject: [PATCH 103/117] fix style

---
 .../sampling/heterogeneous_biased_neighbor_sampling.cpp    | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index ce0b3471cd6..31e8a6b5041 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -252,10 +252,9 @@ INSTANTIATE_TEST_SUITE_P(
   ::testing::Combine(
     ::testing::Values(Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, false},
                       Heterogeneous_Biased_Neighbor_Sampling_Usecase{{4, 10, 7, 8}, 128, 2, true}),
-    ::testing::Values(
-      cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
-      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
-      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
+    ::testing::Values(cugraph::test::File_Usecase("test/datasets/web-Google.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/ljournal-2008.mtx"),
+                      cugraph::test::File_Usecase("test/datasets/webbase-1M.mtx"))));
 
 INSTANTIATE_TEST_SUITE_P(
   rmat_small_test,

From 8cb0c947557cde4294c95ce36f40a616c5f4d41d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Fri, 1 Nov 2024 15:22:38 -0700
Subject: [PATCH 104/117] rename 'd_value' to 'd_span'

---
 .../cugraph/detail/utility_wrappers.hpp       | 18 +++++++--------
 cpp/src/detail/utility_wrappers_32.cu         | 10 ++++----
 cpp/src/detail/utility_wrappers_64.cu         | 10 ++++----
 cpp/src/detail/utility_wrappers_impl.cuh      | 23 ++++++++-----------
 4 files changed, 28 insertions(+), 33 deletions(-)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index 585ee4658f7..c0be99f413c 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -66,46 +66,46 @@ template <typename value_t>
 void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, value_t value);
 
 /**
- * @brief    Sort a buffer
+ * @brief    Sort a device span
  *
  * @tparam      value_t      type of the value to operate on
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[out]  d_value      device array to sort
+ * @param[out]  d_span      device span to sort
  *
  */
 template <typename value_t>
-void sort(raft::handle_t const& handle, raft::device_span<value_t> d_value);
+void sort(raft::handle_t const& handle, raft::device_span<value_t> d_span);
 
 /**
- * @brief    Keep unique element from a buffer
+ * @brief    Keep unique element from a device span
  *
  * @tparam      value_t      type of the value to operate on
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[in]  d_value      device array of unique elements.
+ * @param[in]  d_span      device span of unique elements.
  * @return the number of unique elements.
  *
  */
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value);
+size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_span);
 
 /**
- * @brief    Increment the values of a buffer by a constant value
+ * @brief    Increment the values of a device span by a constant value
  *
  * @tparam      value_t      type of the value to operate on
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[out]  d_value      device array to update
+ * @param[out]  d_span       device span to update
  * @param[in]   value        value to be added to each element of the buffer
  *
  */
 template <typename value_t>
 void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         raft::device_span<value_t> d_value,
+                         raft::device_span<value_t> d_span,
                          value_t value);
 
 /**
diff --git a/cpp/src/detail/utility_wrappers_32.cu b/cpp/src/detail/utility_wrappers_32.cu
index 3fe1a507870..c59006985a9 100644
--- a/cpp/src/detail/utility_wrappers_32.cu
+++ b/cpp/src/detail/utility_wrappers_32.cu
@@ -63,10 +63,10 @@ template void scalar_fill(raft::handle_t const& handle, size_t* d_value, size_t
 
 template void scalar_fill(raft::handle_t const& handle, float* d_value, size_t size, float value);
 
-template void sort(raft::handle_t const& handle, raft::device_span<int32_t> d_value);
+template void sort(raft::handle_t const& handle, raft::device_span<int32_t> d_span);
 
-template size_t unique(raft::handle_t const& handle, raft::device_span<int32_t> d_value);
-template size_t unique(raft::handle_t const& handle, raft::device_span<uint32_t> d_value);
+template size_t unique(raft::handle_t const& handle, raft::device_span<int32_t> d_span);
+template size_t unique(raft::handle_t const& handle, raft::device_span<uint32_t> d_span);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int32_t* d_value,
@@ -79,11 +79,11 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             uint32_t start_value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<int32_t> d_value,
+                                  raft::device_span<int32_t> d_span,
                                   int32_t value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<uint32_t> d_value,
+                                  raft::device_span<uint32_t> d_span,
                                   uint32_t value);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
diff --git a/cpp/src/detail/utility_wrappers_64.cu b/cpp/src/detail/utility_wrappers_64.cu
index 4012aa69b5c..f78c709024a 100644
--- a/cpp/src/detail/utility_wrappers_64.cu
+++ b/cpp/src/detail/utility_wrappers_64.cu
@@ -61,10 +61,10 @@ template void scalar_fill(raft::handle_t const& handle,
 
 template void scalar_fill(raft::handle_t const& handle, double* d_value, size_t size, double value);
 
-template void sort(raft::handle_t const& handle, raft::device_span<int64_t> d_value);
+template void sort(raft::handle_t const& handle, raft::device_span<int64_t> d_span);
 
-template size_t unique(raft::handle_t const& handle, raft::device_span<int64_t> d_value);
-template size_t unique(raft::handle_t const& handle, raft::device_span<uint64_t> d_value);
+template size_t unique(raft::handle_t const& handle, raft::device_span<int64_t> d_span);
+template size_t unique(raft::handle_t const& handle, raft::device_span<uint64_t> d_span);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int64_t* d_value,
@@ -77,11 +77,11 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             uint64_t start_value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<int64_t> d_value,
+                                  raft::device_span<int64_t> d_span,
                                   int64_t value);
 
 template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<uint64_t> d_value,
+                                  raft::device_span<uint64_t> d_span,
                                   uint64_t value);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index 320ab65d37f..7522f3a9c56 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -65,17 +65,17 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
 }
 
 template <typename value_t>
-void sort(raft::handle_t const& handle, raft::device_span<value_t> d_value)
+void sort(raft::handle_t const& handle, raft::device_span<value_t> d_span)
 {
-  thrust::sort(handle.get_thrust_policy(), d_value.begin(), d_value.end());
+  thrust::sort(handle.get_thrust_policy(), d_span.begin(), d_span.end());
 }
 
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_value)
+size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_span)
 {
   auto unique_element_last =
-    thrust::unique(handle.get_thrust_policy(), d_value.begin(), d_value.end());
-  return thrust::distance(d_value.begin(), unique_element_last);
+    thrust::unique(handle.get_thrust_policy(), d_span.begin(), d_span.end());
+  return thrust::distance(d_span.begin(), unique_element_last);
 }
 
 template <typename value_t>
@@ -89,23 +89,18 @@ void sequence_fill(rmm::cuda_stream_view const& stream_view,
 
 template <typename value_t>
 void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         raft::device_span<value_t> d_value,
+                         raft::device_span<value_t> d_span,
                          value_t incr)
 {
   thrust::transform(rmm::exec_policy(stream_view),
-                    d_value.begin(),
-                    d_value.end(),
-                    d_value.begin(),
+                    d_span.begin(),
+                    d_span.end(),
+                    d_span.begin(),
                     cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
                       return static_cast<value_t>(value + incr);
                     }));
 }
 
-template <typename value_t>
-void transform_increment_(rmm::cuda_stream_view const& stream_view, value_t d_value)
-{
-}
-
 template <typename value_t>
 void stride_fill(rmm::cuda_stream_view const& stream_view,
                  value_t* d_value,

From 77303d5be76b89d8392ea746f86e924c4522e375 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 5 Nov 2024 14:59:38 -0800
Subject: [PATCH 105/117] use 'label_list' as a map form 'comm_rank' to
 'label_map'

---
 cpp/src/sampling/detail/conversion_utilities.cu | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index a067005514d..a30f552e6c3 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -45,9 +45,9 @@ rmm::device_uvector<int32_t> flatten_label_map(
 
   thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
   thrust::scatter(handle.get_thrust_policy(),
-                  std::get<0>(label_to_output_comm_rank).begin(),
-                  std::get<0>(label_to_output_comm_rank).end(),
                   std::get<1>(label_to_output_comm_rank).begin(),
+                  std::get<1>(label_to_output_comm_rank).end(),
+                  std::get<0>(label_to_output_comm_rank).begin(),
                   label_map.begin());
 
   return label_map;

From 117a4ecf9a7c1def5500a9d4e3b808aafe42640e Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Tue, 5 Nov 2024 15:02:10 -0800
Subject: [PATCH 106/117] use 'label_list' as a map form 'comm_rank' to
 'label_map'

---
 cpp/src/sampling/detail/conversion_utilities_impl.cuh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
index 6eff0154581..1061b28c44b 100644
--- a/cpp/src/sampling/detail/conversion_utilities_impl.cuh
+++ b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
@@ -46,9 +46,9 @@ rmm::device_uvector<int32_t> flatten_label_map(
 
   thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
   thrust::scatter(handle.get_thrust_policy(),
-                  std::get<0>(label_to_output_comm_rank).begin(),
-                  std::get<0>(label_to_output_comm_rank).end(),
                   std::get<1>(label_to_output_comm_rank).begin(),
+                  std::get<1>(label_to_output_comm_rank).end(),
+                  std::get<0>(label_to_output_comm_rank).begin(),
                   label_map.begin());
 
   return label_map;

From ac66f13216e6dbb7301f490bc7474908f6aa5380 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Wed, 6 Nov 2024 13:04:48 -0800
Subject: [PATCH 107/117] add module biased_neighbor_sample

---
 python/pylibcugraph/pylibcugraph/__init__.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/pylibcugraph/pylibcugraph/__init__.py b/python/pylibcugraph/pylibcugraph/__init__.py
index b9a981094c8..5aa351f9ce1 100644
--- a/python/pylibcugraph/pylibcugraph/__init__.py
+++ b/python/pylibcugraph/pylibcugraph/__init__.py
@@ -40,6 +40,9 @@
 from pylibcugraph.bfs import bfs
 
 from pylibcugraph.uniform_neighbor_sample import uniform_neighbor_sample
+
+from pylibcugraph.biased_neighbor_sample import biased_neighbor_sample
+
 from pylibcugraph.homogeneous_uniform_neighbor_sample import (
     homogeneous_uniform_neighbor_sample,
 )

From 114bf56d68966d9c442c6d2bfe244dc092b6474e Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 12:12:22 -0800
Subject: [PATCH 108/117] rename variable

---
 cpp/include/cugraph/detail/utility_wrappers.hpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index c0be99f413c..681ab319075 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -72,11 +72,11 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[out]  d_span      device span to sort
+ * @param[out]  values      device span to sort
  *
  */
 template <typename value_t>
-void sort(raft::handle_t const& handle, raft::device_span<value_t> d_span);
+void sort(raft::handle_t const& handle, raft::device_span<value_t> values);
 
 /**
  * @brief    Keep unique element from a device span
@@ -85,12 +85,12 @@ void sort(raft::handle_t const& handle, raft::device_span<value_t> d_span);
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[in]  d_span      device span of unique elements.
+ * @param[in]  values      device span of unique elements.
  * @return the number of unique elements.
  *
  */
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_span);
+size_t unique(raft::handle_t const& handle, raft::device_span<value_t> values);
 
 /**
  * @brief    Increment the values of a device span by a constant value
@@ -99,13 +99,13 @@ size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_span);
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
- * @param[out]  d_span       device span to update
+ * @param[out]  values       device span to update
  * @param[in]   value        value to be added to each element of the buffer
  *
  */
 template <typename value_t>
 void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         raft::device_span<value_t> d_span,
+                         raft::device_span<value_t> values,
                          value_t value);
 
 /**

From 63a59ca0f79b65fa7e70acbac855cd6f3fd581e0 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 13:15:12 -0800
Subject: [PATCH 109/117] avoid creating function that compile all types and be
 more explicit when naming

---
 .../cugraph/detail/shuffle_wrappers.hpp       | 24 -------------------
 .../cugraph/detail/utility_wrappers.hpp       | 21 ++++++++--------
 cpp/src/c_api/neighbor_sampling.cpp           | 10 ++++----
 cpp/src/detail/utility_wrappers_32.cu         | 15 ++++--------
 cpp/src/detail/utility_wrappers_64.cu         | 15 ++++--------
 cpp/src/detail/utility_wrappers_impl.cuh      | 22 ++++++++---------
 6 files changed, 36 insertions(+), 71 deletions(-)

diff --git a/cpp/include/cugraph/detail/shuffle_wrappers.hpp b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
index 6592150e748..7dffcce298a 100644
--- a/cpp/include/cugraph/detail/shuffle_wrappers.hpp
+++ b/cpp/include/cugraph/detail/shuffle_wrappers.hpp
@@ -142,30 +142,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<vertex_t>&& vertices,
   rmm::device_uvector<value_t>&& values);
 
-/**
- * @brief Shuffle external (i.e. before renumbering) vertex & values pairs to their local GPU based
- * on vertex partitioning.
- *
- * @tparam vertex_t Type of vertex identifiers. Needs to be an integral type.
- * @tparam value_t Type of values.
- *
- * @param[in] handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
- * and handles to various CUDA libraries) to run graph algorithms.
- * @param[in] vertices Vertices to shuffle.
- * @param[in] values_0 First values to shuffle.
- * @param[in] values_1 Second values to shuffle.
- *
- * @return Tuple of vectors storing shuffled vertex & value pairs.
- */
-template <typename vertex_t, typename value0_t, typename value1_t>
-std::
-  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<vertex_t>&& vertices,
-    rmm::device_uvector<value0_t>&& values_0,
-    rmm::device_uvector<value1_t>&& values_1);
-
 /**
  * @brief Permute a range.
  *
diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index 681ab319075..e0307bc3725 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -68,7 +68,7 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
 /**
  * @brief    Sort a device span
  *
- * @tparam      value_t      type of the value to operate on
+ * @tparam      value_t      type of the value to operate on. Must be either int32_t or int64_t.
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
@@ -76,12 +76,12 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
  *
  */
 template <typename value_t>
-void sort(raft::handle_t const& handle, raft::device_span<value_t> values);
+void sort_ints(raft::handle_t const& handle, raft::device_span<value_t> values);
 
 /**
  * @brief    Keep unique element from a device span
  *
- * @tparam      value_t      type of the value to operate on
+ * @tparam      value_t      type of the value to operate on. Must be either int32_t or int64_t.
  *
  * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
  * and handles to various CUDA libraries) to run graph algorithms.
@@ -90,23 +90,22 @@ void sort(raft::handle_t const& handle, raft::device_span<value_t> values);
  *
  */
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, raft::device_span<value_t> values);
+size_t unique_ints(raft::handle_t const& handle, raft::device_span<value_t> values);
 
 /**
  * @brief    Increment the values of a device span by a constant value
  *
- * @tparam      value_t      type of the value to operate on
+ * @tparam      value_t      type of the value to operate on. Must be either int32_t or int64_t.
  *
- * @param [in]  handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator,
- * and handles to various CUDA libraries) to run graph algorithms.
  * @param[out]  values       device span to update
  * @param[in]   value        value to be added to each element of the buffer
+ * @param[in]   stream_view  stream view
  *
  */
 template <typename value_t>
-void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         raft::device_span<value_t> values,
-                         value_t value);
+void transform_increment_ints(raft::device_span<value_t> values,
+                              value_t value,
+                              rmm::cuda_stream_view const& stream_view);
 
 /**
  * @brief    Fill a buffer with a sequence of values
@@ -116,7 +115,7 @@ void transform_increment(rmm::cuda_stream_view const& stream_view,
  *
  * Similar to the function std::iota, wraps the function thrust::sequence
  *
- * @tparam      value_t      type of the value to operate on
+ * @tparam      value_t      type of the value to operate on. 
  *
  * @param[in]   stream_view  stream view
  * @param[out]  d_value      device array to fill
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 2988c992fc4..842241be81f 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -885,10 +885,10 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
         // Compute the global start_vertex_label_offsets
 
-        cugraph::detail::transform_increment(
-          handle_.get_stream(),
+        cugraph::detail::transform_increment_ints(
           raft::device_span<label_t>{(*start_vertex_labels).data(), (*start_vertex_labels).size()},
-          (label_t)global_labels[handle_.get_comms().get_rank()]);
+          (label_t)global_labels[handle_.get_comms().get_rank()],
+          handle_.get_stream());
       }
 
       if constexpr (multi_gpu) {
@@ -902,11 +902,11 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
 
           // Get unique labels
           // sort the start_vertex_labels
-          cugraph::detail::sort(
+          cugraph::detail::sort_ints(
             handle_.get_stream(),
             raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
 
-          auto num_unique_labels = cugraph::detail::unique(
+          auto num_unique_labels = cugraph::detail::unique_ints(
             handle_.get_stream(),
             raft::device_span<label_t>{unique_labels.data(), unique_labels.size()});
 
diff --git a/cpp/src/detail/utility_wrappers_32.cu b/cpp/src/detail/utility_wrappers_32.cu
index c59006985a9..879a1adf337 100644
--- a/cpp/src/detail/utility_wrappers_32.cu
+++ b/cpp/src/detail/utility_wrappers_32.cu
@@ -63,10 +63,9 @@ template void scalar_fill(raft::handle_t const& handle, size_t* d_value, size_t
 
 template void scalar_fill(raft::handle_t const& handle, float* d_value, size_t size, float value);
 
-template void sort(raft::handle_t const& handle, raft::device_span<int32_t> d_span);
+template void sort_ints(raft::handle_t const& handle, raft::device_span<int32_t> values);
 
-template size_t unique(raft::handle_t const& handle, raft::device_span<int32_t> d_span);
-template size_t unique(raft::handle_t const& handle, raft::device_span<uint32_t> d_span);
+template size_t unique_ints(raft::handle_t const& handle, raft::device_span<int32_t> values);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int32_t* d_value,
@@ -78,13 +77,9 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             size_t size,
                             uint32_t start_value);
 
-template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<int32_t> d_span,
-                                  int32_t value);
-
-template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<uint32_t> d_span,
-                                  uint32_t value);
+template void transform_increment_ints(raft::device_span<int32_t> values,
+                                       int32_t value,
+                                       rmm::cuda_stream_view const& stream_view);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int32_t* d_value,
diff --git a/cpp/src/detail/utility_wrappers_64.cu b/cpp/src/detail/utility_wrappers_64.cu
index f78c709024a..742cb18d718 100644
--- a/cpp/src/detail/utility_wrappers_64.cu
+++ b/cpp/src/detail/utility_wrappers_64.cu
@@ -61,10 +61,9 @@ template void scalar_fill(raft::handle_t const& handle,
 
 template void scalar_fill(raft::handle_t const& handle, double* d_value, size_t size, double value);
 
-template void sort(raft::handle_t const& handle, raft::device_span<int64_t> d_span);
+template void sort_ints(raft::handle_t const& handle, raft::device_span<int64_t> values);
 
-template size_t unique(raft::handle_t const& handle, raft::device_span<int64_t> d_span);
-template size_t unique(raft::handle_t const& handle, raft::device_span<uint64_t> d_span);
+template size_t unique_ints(raft::handle_t const& handle, raft::device_span<int64_t> values);
 
 template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             int64_t* d_value,
@@ -76,13 +75,9 @@ template void sequence_fill(rmm::cuda_stream_view const& stream_view,
                             size_t size,
                             uint64_t start_value);
 
-template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<int64_t> d_span,
-                                  int64_t value);
-
-template void transform_increment(rmm::cuda_stream_view const& stream_view,
-                                  raft::device_span<uint64_t> d_span,
-                                  uint64_t value);
+template void transform_increment_ints(raft::device_span<int64_t> values,
+                                       int64_t value,
+                                       rmm::cuda_stream_view const& stream_view);
 
 template void stride_fill(rmm::cuda_stream_view const& stream_view,
                           int64_t* d_value,
diff --git a/cpp/src/detail/utility_wrappers_impl.cuh b/cpp/src/detail/utility_wrappers_impl.cuh
index 7522f3a9c56..93bd14c4d06 100644
--- a/cpp/src/detail/utility_wrappers_impl.cuh
+++ b/cpp/src/detail/utility_wrappers_impl.cuh
@@ -65,17 +65,17 @@ void scalar_fill(raft::handle_t const& handle, value_t* d_value, size_t size, va
 }
 
 template <typename value_t>
-void sort(raft::handle_t const& handle, raft::device_span<value_t> d_span)
+void sort_ints(raft::handle_t const& handle, raft::device_span<value_t> values)
 {
-  thrust::sort(handle.get_thrust_policy(), d_span.begin(), d_span.end());
+  thrust::sort(handle.get_thrust_policy(), values.begin(), values.end());
 }
 
 template <typename value_t>
-size_t unique(raft::handle_t const& handle, raft::device_span<value_t> d_span)
+size_t unique_ints(raft::handle_t const& handle, raft::device_span<value_t> values)
 {
   auto unique_element_last =
-    thrust::unique(handle.get_thrust_policy(), d_span.begin(), d_span.end());
-  return thrust::distance(d_span.begin(), unique_element_last);
+    thrust::unique(handle.get_thrust_policy(), values.begin(), values.end());
+  return thrust::distance(values.begin(), unique_element_last);
 }
 
 template <typename value_t>
@@ -88,14 +88,14 @@ void sequence_fill(rmm::cuda_stream_view const& stream_view,
 }
 
 template <typename value_t>
-void transform_increment(rmm::cuda_stream_view const& stream_view,
-                         raft::device_span<value_t> d_span,
-                         value_t incr)
+void transform_increment_ints(raft::device_span<value_t> values,
+                              value_t incr,
+                              rmm::cuda_stream_view const& stream_view)
 {
   thrust::transform(rmm::exec_policy(stream_view),
-                    d_span.begin(),
-                    d_span.end(),
-                    d_span.begin(),
+                    values.begin(),
+                    values.end(),
+                    values.begin(),
                     cuda::proclaim_return_type<value_t>([incr] __device__(value_t value) {
                       return static_cast<value_t>(value + incr);
                     }));

From 84face3bb1f386bfa2cf4f4f3e08c2a62a55288b Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 18:31:37 -0800
Subject: [PATCH 110/117] remove unsued function

---
 cpp/src/utilities/shuffle_vertices.cuh        | 30 -------------------
 .../utilities/shuffle_vertices_mg_v32_fp.cu   | 16 ----------
 .../shuffle_vertices_mg_v32_integral.cu       | 16 ----------
 .../utilities/shuffle_vertices_mg_v64_fp.cu   | 16 ----------
 .../shuffle_vertices_mg_v64_integral.cu       | 16 ----------
 5 files changed, 94 deletions(-)

diff --git a/cpp/src/utilities/shuffle_vertices.cuh b/cpp/src/utilities/shuffle_vertices.cuh
index 59b5ecad059..5ed6513816f 100644
--- a/cpp/src/utilities/shuffle_vertices.cuh
+++ b/cpp/src/utilities/shuffle_vertices.cuh
@@ -131,36 +131,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   return std::make_tuple(std::move(d_vertices), std::move(d_values));
 }
 
-template <typename vertex_t, typename value0_t, typename value1_t>
-std::
-  tuple<rmm::device_uvector<vertex_t>, rmm::device_uvector<value0_t>, rmm::device_uvector<value1_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<vertex_t>&& vertices,
-    rmm::device_uvector<value0_t>&& values_0,
-    rmm::device_uvector<value1_t>&& values_1)
-{
-  auto const comm_size       = handle.get_comms().get_size();
-  auto& major_comm           = handle.get_subcomm(cugraph::partition_manager::major_comm_name());
-  auto const major_comm_size = major_comm.get_size();
-  auto& minor_comm           = handle.get_subcomm(cugraph::partition_manager::minor_comm_name());
-  auto const minor_comm_size = minor_comm.get_size();
-
-  rmm::device_uvector<vertex_t> d_vertices(0, handle.get_stream());
-  rmm::device_uvector<value0_t> d_values_0(0, handle.get_stream());
-  std::optional<rmm::device_uvector<value1_t>> d_values_1(std::nullopt);
-
-  std::tie(d_vertices, d_values_0, d_values_1) = shuffle_vertices_and_values_by_gpu_id_impl(
-    handle,
-    std::move(vertices),
-    std::move(values_0),
-    std::make_optional(std::move(values_1)),
-    cugraph::detail::compute_gpu_id_from_ext_vertex_t<vertex_t>{
-      comm_size, major_comm_size, minor_comm_size});
-
-  return std::make_tuple(std::move(d_vertices), std::move(d_values_0), std::move(*d_values_1));
-}
-
 template <typename vertex_t>
 rmm::device_uvector<vertex_t> shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
   raft::handle_t const& handle,
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
index c66a0868b6f..1112d7af2f7 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_fp.cu
@@ -40,22 +40,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int32_t>&& vertices,
-    rmm::device_uvector<float>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int32_t>&& vertices,
-    rmm::device_uvector<double>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<float>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
index 984f2afc002..db7be5a3031 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v32_integral.cu
@@ -62,22 +62,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int32_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int32_t>&& vertices,
-    rmm::device_uvector<int32_t>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
-template std::
-  tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int32_t>&& vertices,
-    rmm::device_uvector<size_t>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
index ff200fac563..9592a9c63d1 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_fp.cu
@@ -40,22 +40,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<double>&& values);
 
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int64_t>&& vertices,
-    rmm::device_uvector<float>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<double>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int64_t>&& vertices,
-    rmm::device_uvector<double>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
 }  // namespace detail
 
 template std::tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<float>>
diff --git a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
index fb4c0d469f7..7d968006bc7 100644
--- a/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
+++ b/cpp/src/utilities/shuffle_vertices_mg_v64_integral.cu
@@ -54,22 +54,6 @@ shuffle_ext_vertex_value_pairs_to_local_gpu_by_vertex_partitioning(
   rmm::device_uvector<int64_t>&& vertices,
   rmm::device_uvector<size_t>&& values);
 
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<int32_t>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int64_t>&& vertices,
-    rmm::device_uvector<int32_t>&& values_0,
-    rmm::device_uvector<int32_t>&& values_1);
-
-template std::
-  tuple<rmm::device_uvector<int64_t>, rmm::device_uvector<size_t>, rmm::device_uvector<int32_t>>
-  shuffle_ext_vertex_values_pairs_to_local_gpu_by_vertex_partitioning(
-    raft::handle_t const& handle,
-    rmm::device_uvector<int64_t>&& vertices,
-    rmm::device_uvector<size_t>&& values,
-    rmm::device_uvector<int32_t>&& values_1);
-
 template rmm::device_uvector<int64_t> shuffle_int_vertices_to_local_gpu_by_vertex_partitioning(
   raft::handle_t const& handle,
   rmm::device_uvector<int64_t>&& vertices,

From ab853e54cf975f1e81acff8dd376d656ac69029d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 18:37:13 -0800
Subject: [PATCH 111/117] declare homogeneous functions first

---
 cpp/include/cugraph_c/sampling_algorithms.h | 32 ++++++++++-----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/cpp/include/cugraph_c/sampling_algorithms.h b/cpp/include/cugraph_c/sampling_algorithms.h
index dde2e270441..b6d51165da2 100644
--- a/cpp/include/cugraph_c/sampling_algorithms.h
+++ b/cpp/include/cugraph_c/sampling_algorithms.h
@@ -450,7 +450,7 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
   cugraph_error_t** error);
 
 /**
- * @brief     Heterogeneous Uniform Neighborhood Sampling
+ * @brief     Homogeneous Uniform Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices and heterogeneous
  * fan_out types. The neighborhood is sampled uniformly.
@@ -470,8 +470,6 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  * This parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
- * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
- * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -481,24 +479,23 @@ cugraph_error_code_t cugraph_biased_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
+cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
 /**
- * @brief     Heterogeneous Biased Neighborhood Sampling
+ * @brief     Homogeneous Biased Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices and heterogeneous
- * fan_out types. The neighborhood is sampled with biases.
+ * fan_out types. The neighborhood is sampled uniformly.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -517,8 +514,6 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  * This parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
- * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
- * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -528,7 +523,7 @@ cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
+cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -536,14 +531,13 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
-  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
 /**
- * @brief     Homogeneous Uniform Neighborhood Sampling
+ * @brief     Heterogeneous Uniform Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices and heterogeneous
  * fan_out types. The neighborhood is sampled uniformly.
@@ -563,6 +557,8 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
  * This parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
+ * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -572,23 +568,24 @@ cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_uniform_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
+  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,
   cugraph_error_t** error);
 
 /**
- * @brief     Homogeneous Biased Neighborhood Sampling
+ * @brief     Heterogeneous Biased Neighborhood Sampling
  *
  * Returns a sample of the neighborhood around specified start vertices and heterogeneous
- * fan_out types. The neighborhood is sampled uniformly.
+ * fan_out types. The neighborhood is sampled with biases.
  * Optionally, each start vertex can be associated with a label, allowing the caller to specify
  * multiple batches of sampling requests in the same function call - which should improve GPU
  * utilization.
@@ -607,6 +604,8 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
  * This parameter is only used with the retain_seeds option.
  * @param [in]  fan_out       Host array defining the fan out at each step in the sampling
  * algorithm. We only support fan_out values of type INT32
+ * @param [in]  num_edge_types Number of edge types where a value of 1 translates to homogeneous
+ * neighbor sample whereas a value greater than 1 translates to heterogeneous neighbor sample.
  * @param [in]  sampling_options
  *                           Opaque pointer defining the sampling options.
  * @param [in]  do_expensive_check
@@ -616,7 +615,7 @@ cugraph_error_code_t cugraph_homogeneous_uniform_neighbor_sample(
  *                           be populated if error code is not CUGRAPH_SUCCESS
  * @return error code
  */
-cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
+cugraph_error_code_t cugraph_heterogeneous_biased_neighbor_sample(
   const cugraph_resource_handle_t* handle,
   cugraph_rng_state_t* rng_state,
   cugraph_graph_t* graph,
@@ -624,6 +623,7 @@ cugraph_error_code_t cugraph_homogeneous_biased_neighbor_sample(
   const cugraph_type_erased_device_array_view_t* start_vertices,
   const cugraph_type_erased_device_array_view_t* start_vertex_offsets,
   const cugraph_type_erased_host_array_view_t* fan_out,
+  int num_edge_types,
   const cugraph_sampling_options_t* options,
   bool_t do_expensive_check,
   cugraph_sample_result_t** result,

From 802d9b03cf033ae1b9e1e22e37f62e368a100dc2 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 18:41:37 -0800
Subject: [PATCH 112/117] rename variable

---
 cpp/include/cugraph/sampling_functions.hpp       | 16 ++++++++--------
 cpp/src/c_api/neighbor_sampling.cpp              |  2 +-
 cpp/src/sampling/detail/conversion_utilities.cu  |  6 +++---
 .../detail/conversion_utilities_impl.cuh         |  6 +++---
 cpp/src/sampling/detail/sampling_utils.hpp       |  6 +++---
 5 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index f3f9b126513..f4178efc83d 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -284,10 +284,10 @@ struct sampling_flags_t {
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
  * offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
  * in the return will be a CSR-style offsets array to identify the beginning of each label range in
  * the output vectors.
  *
@@ -369,10 +369,10 @@ heterogeneous_uniform_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
  * offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
  * in the return will be a CSR-style offsets array to identify the beginning of each label range in
  * the output vectors.
  *
@@ -459,10 +459,10 @@ heterogeneous_biased_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
  * offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
  * in the return will be a CSR-style offsets array to identify the beginning of each label range in
  * the output vectors.
  *
@@ -540,10 +540,10 @@ homogeneous_uniform_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_offsets is not specified then no organization is applied to the output, the
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
  * offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_offsets is specified the offsets array will be populated. The offsets array
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
  * in the return will be a CSR-style offsets array to identify the beginning of each label range in
  * the output vectors.
  *
diff --git a/cpp/src/c_api/neighbor_sampling.cpp b/cpp/src/c_api/neighbor_sampling.cpp
index 842241be81f..f4f545597d4 100644
--- a/cpp/src/c_api/neighbor_sampling.cpp
+++ b/cpp/src/c_api/neighbor_sampling.cpp
@@ -878,7 +878,7 @@ struct neighbor_sampling_functor : public cugraph::c_api::abstract_functor {
           global_labels.begin(), global_labels.end(), global_labels.begin(), label_t{0});
 
         // Retrieve the start_vertex_labels
-        start_vertex_labels = cugraph::detail::convert_starting_vertex_offsets_to_labels(
+        start_vertex_labels = cugraph::detail::convert_starting_vertex_label_offsets_to_labels(
           handle_,
           raft::device_span<size_t const>{start_vertex_offsets_->as_type<size_t>(),
                                           start_vertex_offsets_->size_});
diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index a30f552e6c3..55658f7987a 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -21,10 +21,10 @@
 namespace cugraph {
 namespace detail {
 
-rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
-  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets)
+rmm::device_uvector<int32_t> convert_starting_vertex_label_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_label_offsets)
 {
-  return expand_sparse_offsets(starting_vertex_offsets, int32_t{0}, handle.get_stream());
+  return expand_sparse_offsets(starting_vertex_label_offsets, int32_t{0}, handle.get_stream());
 }
 
 template <typename label_t>
diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
index 1061b28c44b..d936abac0ae 100644
--- a/cpp/src/sampling/detail/conversion_utilities_impl.cuh
+++ b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
@@ -23,10 +23,10 @@
 namespace cugraph {
 namespace detail {
 
-rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
-  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets)
+rmm::device_uvector<int32_t> convert_starting_vertex_label_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_label_offsets)
 {
-  return expand_sparse_offsets(*starting_vertex_offsets, int32_t{0}, handle.get_stream());
+  return expand_sparse_offsets(*starting_vertex_label_offsets, int32_t{0}, handle.get_stream());
 }
 
 rmm::device_uvector<int32_t> flatten_label_map(
diff --git a/cpp/src/sampling/detail/sampling_utils.hpp b/cpp/src/sampling/detail/sampling_utils.hpp
index b68f3fc74f7..17eb8dd0873 100644
--- a/cpp/src/sampling/detail/sampling_utils.hpp
+++ b/cpp/src/sampling/detail/sampling_utils.hpp
@@ -300,12 +300,12 @@ shuffle_and_organize_output(
  *
  * @param handle RAFT handle object to encapsulate resources (e.g. CUDA stream, communicator, and
  * handles to various CUDA libraries) to run graph algorithms.
- * @param starting_vertex_offsets Offsets array defining where each vertex label begins
+ * @param starting_vertex_label_offsets Offsets array defining where each vertex label begins
  *
  * @returns device vector containing labels for each starting vertex
  */
-rmm::device_uvector<int32_t> convert_starting_vertex_offsets_to_labels(
-  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_offsets);
+rmm::device_uvector<int32_t> convert_starting_vertex_label_offsets_to_labels(
+  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_label_offsets);
 
 /**
  * @brief   Flatten the legacy label_to_output_comm_rank into the new structure

From c5bec5f2358ddc48e73a1ad1e57f05afdd802665 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 19:38:46 -0800
Subject: [PATCH 113/117] remove duplicated functions

---
 .../detail/conversion_utilities_impl.cuh      | 58 -------------------
 1 file changed, 58 deletions(-)
 delete mode 100644 cpp/src/sampling/detail/conversion_utilities_impl.cuh

diff --git a/cpp/src/sampling/detail/conversion_utilities_impl.cuh b/cpp/src/sampling/detail/conversion_utilities_impl.cuh
deleted file mode 100644
index d936abac0ae..00000000000
--- a/cpp/src/sampling/detail/conversion_utilities_impl.cuh
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * Copyright (c) 2024, NVIDIA CORPORATION.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include "sampling/detail/sampling_utils.hpp"
-
-#include <cugraph/utilities/misc_utils.cuh>
-
-namespace cugraph {
-namespace detail {
-
-rmm::device_uvector<int32_t> convert_starting_vertex_label_offsets_to_labels(
-  raft::handle_t const& handle, raft::device_span<size_t const> starting_vertex_label_offsets)
-{
-  return expand_sparse_offsets(*starting_vertex_label_offsets, int32_t{0}, handle.get_stream());
-}
-
-rmm::device_uvector<int32_t> flatten_label_map(
-  raft::handle_t const& handle,
-  std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>
-    label_to_output_comm_rank)
-{
-  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
-
-  label_t max_label = thrust::reduce(handle.get_thrust_policy(),
-                                     std::get<0>(label_to_output_comm_rank).begin(),
-                                     std::get<0>(label_to_output_comm_rank).end(),
-                                     label_t{0},
-                                     thrust::maximum<label_t>());
-
-  label_map.resize(max_label + 1, handle.get_stream());
-
-  thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
-  thrust::scatter(handle.get_thrust_policy(),
-                  std::get<1>(label_to_output_comm_rank).begin(),
-                  std::get<1>(label_to_output_comm_rank).end(),
-                  std::get<0>(label_to_output_comm_rank).begin(),
-                  label_map.begin());
-
-  return label_map;
-}
-
-}  // namespace detail
-}  // namespace cugraph

From 4bd09f6e22e7505880012c58518b5a36d3049218 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 20:05:11 -0800
Subject: [PATCH 114/117] reorder variable declaration

---
 cpp/src/sampling/detail/conversion_utilities.cu | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/cpp/src/sampling/detail/conversion_utilities.cu b/cpp/src/sampling/detail/conversion_utilities.cu
index 55658f7987a..0279735dc1f 100644
--- a/cpp/src/sampling/detail/conversion_utilities.cu
+++ b/cpp/src/sampling/detail/conversion_utilities.cu
@@ -33,15 +33,13 @@ rmm::device_uvector<int32_t> flatten_label_map(
   std::tuple<raft::device_span<label_t const>, raft::device_span<int32_t const>>
     label_to_output_comm_rank)
 {
-  rmm::device_uvector<int32_t> label_map(0, handle.get_stream());
-
   label_t max_label = thrust::reduce(handle.get_thrust_policy(),
                                      std::get<0>(label_to_output_comm_rank).begin(),
                                      std::get<0>(label_to_output_comm_rank).end(),
                                      label_t{0},
                                      thrust::maximum<label_t>());
 
-  label_map.resize(max_label + 1, handle.get_stream());
+  rmm::device_uvector<int32_t> label_map(max_label + 1, handle.get_stream());
 
   thrust::fill(handle.get_thrust_policy(), label_map.begin(), label_map.end(), int32_t{0});
   thrust::scatter(handle.get_thrust_policy(),

From 5d6cb347cdd8cc0711870079904ce3cfad9677ad Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 20:16:10 -0800
Subject: [PATCH 115/117] add fixme for not testing edge masking

---
 .../sampling/heterogeneous_biased_neighbor_sampling.cpp    | 7 ++++++-
 .../sampling/heterogeneous_uniform_neighbor_sampling.cpp   | 7 ++++++-
 .../sampling/mg_heterogeneous_biased_neighbor_sampling.cpp | 7 ++++++-
 .../mg_heterogeneous_uniform_neighbor_sampling.cpp         | 7 ++++++-
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index 31e8a6b5041..f6807ecfc8c 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -28,7 +28,12 @@ struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
   int32_t batch_size{10};
   int32_t num_edge_types{1};
   bool flag_replacement{true};
-
+  // FIXME: Edge masking is currently not tested because it will
+  // require attaching two masks (edge_type_t, bool_t) which
+  // is not currently supported. Once a primitive to support 
+  // heterogeneous sampling is added, maintaining two masks
+  // won't be necessary
+  // bool edge_masking{false};
   bool check_correctness{true};
 };
 
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index db2e3fbcb05..f88a80419d8 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -28,7 +28,12 @@ struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
   int32_t batch_size{10};
   int32_t num_edge_types{1};
   bool flag_replacement{true};
-
+  // FIXME: Edge masking is currently not tested because it will
+  // require attaching two masks (edge_type_t, bool_t) which
+  // is not currently supported. Once a primitive to support 
+  // heterogeneous sampling is added, maintaining two masks
+  // won't be necessary
+  // bool edge_masking{false};
   bool check_correctness{true};
 };
 
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index e08dcb087cc..fd70aacf329 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -31,7 +31,12 @@ struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
   int32_t batch_size{10};
   int32_t num_edge_types{1};
   bool with_replacement{true};
-
+  // FIXME: Edge masking is currently not tested because it will
+  // require attaching two masks (edge_type_t, bool_t) which
+  // is not currently supported. Once a primitive to support 
+  // heterogeneous sampling is added, maintaining two masks
+  // won't be necessary
+  // bool edge_masking{false};
   bool check_correctness{true};
 };
 
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index bf338b964f0..eed1c4078f0 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -31,7 +31,12 @@ struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
   int32_t batch_size{10};
   int32_t num_edge_types{1};
   bool with_replacement{true};
-
+  // FIXME: Edge masking is currently not tested because it will
+  // require attaching two masks (edge_type_t, bool_t) which
+  // is not currently supported. Once a primitive to support 
+  // heterogeneous sampling is added, maintaining two masks
+  // won't be necessary
+  // bool edge_masking{false};
   bool check_correctness{true};
 };
 

From 24e31cb3eec92c5b611673bc8ab7baabbfe1bc1d Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 20:18:22 -0800
Subject: [PATCH 116/117] remove outdated fixme

---
 cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp | 2 --
 cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp    | 2 --
 2 files changed, 4 deletions(-)

diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index f88a80419d8..762381b60f8 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -82,8 +82,6 @@ class Tests_Heterogeneous_Uniform_Neighbor_Sampling
 
     constexpr float select_probability{0.05};
 
-    // FIXME:  Update the tests to initialize RngState and use it instead
-    //         of seed...
     constexpr uint64_t seed{0};
 
     raft::random::RngState rng_state(seed);
diff --git a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
index 3e4d8ea2adc..14cf54e7d1c 100644
--- a/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/homogeneous_biased_neighbor_sampling.cpp
@@ -82,8 +82,6 @@ class Tests_Homogeneous_Biased_Neighbor_Sampling
 
     constexpr float select_probability{0.05};
 
-    // FIXME:  Update the tests to initialize RngState and use it instead
-    //         of seed...
     constexpr uint64_t seed{0};
 
     raft::random::RngState rng_state(seed);

From 7ca5d5926f632e444333be1ec9217878fb3440d3 Mon Sep 17 00:00:00 2001
From: jnke2016 <jnke2016@gmail.com>
Date: Thu, 7 Nov 2024 20:21:22 -0800
Subject: [PATCH 117/117] fix style

---
 .../cugraph/detail/utility_wrappers.hpp       |  2 +-
 cpp/include/cugraph/sampling_functions.hpp    | 40 +++++++++----------
 ...heterogeneous_biased_neighbor_sampling.cpp |  2 +-
 ...eterogeneous_uniform_neighbor_sampling.cpp |  2 +-
 ...heterogeneous_biased_neighbor_sampling.cpp |  2 +-
 ...eterogeneous_uniform_neighbor_sampling.cpp |  2 +-
 6 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/cpp/include/cugraph/detail/utility_wrappers.hpp b/cpp/include/cugraph/detail/utility_wrappers.hpp
index e0307bc3725..b1afeafd66b 100644
--- a/cpp/include/cugraph/detail/utility_wrappers.hpp
+++ b/cpp/include/cugraph/detail/utility_wrappers.hpp
@@ -115,7 +115,7 @@ void transform_increment_ints(raft::device_span<value_t> values,
  *
  * Similar to the function std::iota, wraps the function thrust::sequence
  *
- * @tparam      value_t      type of the value to operate on. 
+ * @tparam      value_t      type of the value to operate on.
  *
  * @param[in]   stream_view  stream view
  * @param[out]  d_value      device array to fill
diff --git a/cpp/include/cugraph/sampling_functions.hpp b/cpp/include/cugraph/sampling_functions.hpp
index f4178efc83d..5033e46f339 100644
--- a/cpp/include/cugraph/sampling_functions.hpp
+++ b/cpp/include/cugraph/sampling_functions.hpp
@@ -284,12 +284,12 @@ struct sampling_flags_t {
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
- * offsets values in the return set will be std::nullopt.
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the
+ * output, the offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
- * in the return will be a CSR-style offsets array to identify the beginning of each label range in
- * the output vectors.
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets
+ * array in the return will be a CSR-style offsets array to identify the beginning of each label
+ * range in the output vectors.
  *
  * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
  * for a particular label are returned on the specified rank.  This will result in the offsets array
@@ -369,12 +369,12 @@ heterogeneous_uniform_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
- * offsets values in the return set will be std::nullopt.
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the
+ * output, the offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
- * in the return will be a CSR-style offsets array to identify the beginning of each label range in
- * the output vectors.
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets
+ * array in the return will be a CSR-style offsets array to identify the beginning of each label
+ * range in the output vectors.
  *
  * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
  * for a particular label are returned on the specified rank.  This will result in the offsets array
@@ -459,12 +459,12 @@ heterogeneous_biased_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
- * offsets values in the return set will be std::nullopt.
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the
+ * output, the offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
- * in the return will be a CSR-style offsets array to identify the beginning of each label range in
- * the output vectors.
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets
+ * array in the return will be a CSR-style offsets array to identify the beginning of each label
+ * range in the output vectors.
  *
  * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
  * for a particular label are returned on the specified rank.  This will result in the offsets array
@@ -540,12 +540,12 @@ homogeneous_uniform_neighbor_sample(
  * id, edge_type (optional) identifies the edge type, hop identifies which hop the edge was
  * encountered in.  The offsets array (optional) identifies the offset for each label.
  *
- * If @p starting_vertex_label_offsets is not specified then no organization is applied to the output, the
- * offsets values in the return set will be std::nullopt.
+ * If @p starting_vertex_label_offsets is not specified then no organization is applied to the
+ * output, the offsets values in the return set will be std::nullopt.
  *
- * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets array
- * in the return will be a CSR-style offsets array to identify the beginning of each label range in
- * the output vectors.
+ * If @p starting_vertex_label_offsets is specified the offsets array will be populated. The offsets
+ * array in the return will be a CSR-style offsets array to identify the beginning of each label
+ * range in the output vectors.
  *
  * If @p label_to_output_comm_rank is specified then the data will be shuffled so that all entries
  * for a particular label are returned on the specified rank.  This will result in the offsets array
diff --git a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
index f6807ecfc8c..6ea00cf5104 100644
--- a/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_biased_neighbor_sampling.cpp
@@ -30,7 +30,7 @@ struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
   bool flag_replacement{true};
   // FIXME: Edge masking is currently not tested because it will
   // require attaching two masks (edge_type_t, bool_t) which
-  // is not currently supported. Once a primitive to support 
+  // is not currently supported. Once a primitive to support
   // heterogeneous sampling is added, maintaining two masks
   // won't be necessary
   // bool edge_masking{false};
diff --git a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
index 762381b60f8..3b57aed4768 100644
--- a/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/heterogeneous_uniform_neighbor_sampling.cpp
@@ -30,7 +30,7 @@ struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
   bool flag_replacement{true};
   // FIXME: Edge masking is currently not tested because it will
   // require attaching two masks (edge_type_t, bool_t) which
-  // is not currently supported. Once a primitive to support 
+  // is not currently supported. Once a primitive to support
   // heterogeneous sampling is added, maintaining two masks
   // won't be necessary
   // bool edge_masking{false};
diff --git a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
index fd70aacf329..abdf8824966 100644
--- a/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_biased_neighbor_sampling.cpp
@@ -33,7 +33,7 @@ struct Heterogeneous_Biased_Neighbor_Sampling_Usecase {
   bool with_replacement{true};
   // FIXME: Edge masking is currently not tested because it will
   // require attaching two masks (edge_type_t, bool_t) which
-  // is not currently supported. Once a primitive to support 
+  // is not currently supported. Once a primitive to support
   // heterogeneous sampling is added, maintaining two masks
   // won't be necessary
   // bool edge_masking{false};
diff --git a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
index eed1c4078f0..42abf18682e 100644
--- a/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
+++ b/cpp/tests/sampling/mg_heterogeneous_uniform_neighbor_sampling.cpp
@@ -33,7 +33,7 @@ struct Heterogeneous_Uniform_Neighbor_Sampling_Usecase {
   bool with_replacement{true};
   // FIXME: Edge masking is currently not tested because it will
   // require attaching two masks (edge_type_t, bool_t) which
-  // is not currently supported. Once a primitive to support 
+  // is not currently supported. Once a primitive to support
   // heterogeneous sampling is added, maintaining two masks
   // won't be necessary
   // bool edge_masking{false};