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 01/68] 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 02/68] 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 03/68] 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 04/68] 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 05/68] 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 06/68] 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 07/68] 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 08/68] 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 09/68] 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 10/68] 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 11/68] 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 12/68] 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 13/68] 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 14/68] 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 15/68] 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 16/68] 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 17/68] 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 18/68] 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 19/68] 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 20/68] 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 21/68] 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 22/68] 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 23/68] 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 24/68] 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 25/68] 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 26/68] 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 27/68] 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 28/68] 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 29/68] 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 30/68] 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 31/68] 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 32/68] 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 33/68] 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 34/68] 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 35/68] 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 36/68] 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 37/68] 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 38/68] 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 39/68] 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 40/68] 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 41/68] 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 42/68] 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 43/68] 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 44/68] 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 45/68] 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 46/68] 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 47/68] 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 48/68] 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 49/68] 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 50/68] 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 51/68] 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 52/68] 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 53/68] 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 54/68] 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 55/68] 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 56/68] 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 57/68] 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 58/68] 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 59/68] 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 60/68] 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 61/68] 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 62/68] 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 63/68] 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 64/68] 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 65/68] 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 66/68] 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 67/68] 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 68/68] 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)