From 242d72529deab50c12816a0619da6d028c2bc15a Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 29 May 2024 16:17:30 +0000
Subject: [PATCH 01/25] enable nn descent dist return

---
 .../raft/neighbors/detail/nn_descent.cuh      | 24 ++++++++++++++++---
 .../raft/neighbors/nn_descent_types.hpp       | 10 ++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index cd3c6f3947..a2903532e8 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -345,7 +345,7 @@ class GNND {
   GNND(const GNND&)            = delete;
   GNND& operator=(const GNND&) = delete;
 
-  void build(Data_t* data, const Index_t nrow, Index_t* output_graph);
+  void build(Data_t* data, const Index_t nrow, Index_t* output_graph, bool return_distances, DistData_t *output_distances);
   ~GNND()    = default;
   using ID_t = InternalID_t<Index_t>;
 
@@ -1212,7 +1212,7 @@ void GNND<Data_t, Index_t>::local_join(cudaStream_t stream)
 }
 
 template <typename Data_t, typename Index_t>
-void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* output_graph)
+void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* output_graph, bool return_distances, DistData_t *output_distances)
 {
   using input_t = typename std::remove_const<Data_t>::type;
 
@@ -1338,6 +1338,14 @@ void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* out
 
   // Reuse graph_.h_dists as the buffer for shrink the lists in graph
   static_assert(sizeof(decltype(*(graph_.h_dists.data_handle()))) >= sizeof(Index_t));
+
+  if(return_distances) {
+    raft::copy(output_distances,
+               graph_.h_dists.data_handle(),
+               nrow_ * build_config_.node_degree,
+               raft::resource::get_cuda_stream(res));
+  }
+
   Index_t* graph_shrink_buffer = (Index_t*)graph_.h_dists.data_handle();
 
 #pragma omp parallel for
@@ -1404,6 +1412,12 @@ void build(raft::resources const& res,
 
   auto int_graph = raft::make_host_matrix<int, int64_t, row_major>(
     dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
+  
+  auto distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(0,0);
+  if(params.return_distances) {
+    distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(
+      dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
+  }
 
   BuildConfig build_config{.max_dataset_size      = static_cast<size_t>(dataset.extent(0)),
                            .dataset_dim           = static_cast<size_t>(dataset.extent(1)),
@@ -1413,13 +1427,17 @@ void build(raft::resources const& res,
                            .termination_threshold = params.termination_threshold};
 
   GNND<const T, int> nnd(res, build_config);
-  nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle());
+  nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle(), params.return_distances, distances_graph.data_handle());
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
     for (size_t j = 0; j < graph_degree; j++) {
       auto graph                  = idx.graph().data_handle();
       graph[i * graph_degree + j] = int_graph.data_handle()[i * extended_graph_degree + j];
+      if(params.return_distances) {
+        auto dist_graph = idx.distances().data_handle();
+        dist_graph[i * graph_degree + j] = distances_graph.data_handle()[i * extended_graph_degree + j];
+      }
     }
   }
 }
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index e1fc96878a..20be0ff749 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -26,6 +26,7 @@
 #include <raft/distance/distance_types.hpp>
 
 namespace raft::neighbors::experimental::nn_descent {
+using DistData_t = float;
 /**
  * @ingroup nn-descent
  * @{
@@ -51,6 +52,7 @@ struct index_params : ann::index_params {
   size_t intermediate_graph_degree = 128;     // Degree of input graph for pruning.
   size_t max_iterations            = 20;      // Number of nn-descent iterations.
   float termination_threshold      = 0.0001;  // Termination threshold of nn-descent.
+  bool return_distances            = 0;       // return distances if 1
 };
 
 /**
@@ -85,6 +87,7 @@ struct index : ann::index {
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(n_rows, n_cols)},
+      distances_{raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols)},
       graph_view_{graph_.view()}
   {
   }
@@ -105,6 +108,7 @@ struct index : ann::index {
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
+      distances_{raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0)},
       graph_view_{graph_view}
   {
   }
@@ -133,6 +137,11 @@ struct index : ann::index {
     return graph_view_;
   }
 
+  [[nodiscard]] inline auto distances() noexcept -> host_matrix_view<DistData_t, int64_t, row_major>
+  {
+    return distances_.view();
+  }
+
   // Don't allow copying the index for performance reasons (try avoiding copying data)
   index(const index&)                    = delete;
   index(index&&)                         = default;
@@ -144,6 +153,7 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
+  raft::host_matrix<DistData_t, int64_t, row_major> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
 };

From 0977947e2a6b146ca9120981eb3c9f0fc70bcb6b Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 29 May 2024 21:41:09 +0000
Subject: [PATCH 02/25] change bool to int

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 20be0ff749..626adcaa90 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -52,7 +52,7 @@ struct index_params : ann::index_params {
   size_t intermediate_graph_degree = 128;     // Degree of input graph for pruning.
   size_t max_iterations            = 20;      // Number of nn-descent iterations.
   float termination_threshold      = 0.0001;  // Termination threshold of nn-descent.
-  bool return_distances            = 0;       // return distances if 1
+  int return_distances             = 0;       // return distances if 1
 };
 
 /**

From a7749acc3a620fb78d8885ea68917bb420cf0d6b Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 30 May 2024 20:38:39 +0000
Subject: [PATCH 03/25] add test for distance

---
 cpp/test/neighbors/ann_nn_descent.cuh | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index 495af081f1..61c8f88c26 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -65,7 +65,9 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
   {
     size_t queries_size = ps.n_rows * ps.graph_degree;
     std::vector<IdxT> indices_NNDescent(queries_size);
+    std::vector<DistanceT> distances_NNDescent(queries_size);
     std::vector<IdxT> indices_naive(queries_size);
+    std::vector<DistanceT> distances_naive(queries_size);
 
     {
       rmm::device_uvector<DistanceT> distances_naive_dev(queries_size, stream_);
@@ -81,6 +83,7 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
                                         ps.graph_degree,
                                         ps.metric);
       update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_);
+      update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_);
       resource::sync_stream(handle_);
     }
 
@@ -104,10 +107,14 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_host_view);
             update_host(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            update_host(
+              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           } else {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_view);
             update_host(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            update_host(
+              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           };
         }
         resource::sync_stream(handle_);
@@ -116,6 +123,8 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
       double min_recall = ps.min_recall;
       EXPECT_TRUE(eval_recall(
         indices_naive, indices_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall));
+      EXPECT_TRUE(eval_recall(
+        distances_naive, distances_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall, false));
     }
   }
 

From a753e45edd78466de1ab41e82598130bbb71d0ad Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 30 May 2024 21:21:01 +0000
Subject: [PATCH 04/25] test for indices and distances with one func

---
 cpp/test/neighbors/ann_nn_descent.cuh | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index 61c8f88c26..36ac32eb24 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -121,10 +121,8 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
       }
 
       double min_recall = ps.min_recall;
-      EXPECT_TRUE(eval_recall(
-        indices_naive, indices_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall));
-      EXPECT_TRUE(eval_recall(
-        distances_naive, distances_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall, false));
+      EXPECT_TRUE(eval_neighbours(
+        indices_naive, indices_NNDescent, distances_naive, distances_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall));
     }
   }
 

From 1d96e1761aff5aeb25c571e368c7ce97526e5da1 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 30 May 2024 21:28:47 +0000
Subject: [PATCH 05/25] fix styling

---
 .../raft/neighbors/detail/nn_descent.cuh      | 31 +++++++++++++------
 cpp/test/neighbors/ann_nn_descent.cuh         | 10 ++++--
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index a2903532e8..c617c05c6a 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -345,7 +345,11 @@ class GNND {
   GNND(const GNND&)            = delete;
   GNND& operator=(const GNND&) = delete;
 
-  void build(Data_t* data, const Index_t nrow, Index_t* output_graph, bool return_distances, DistData_t *output_distances);
+  void build(Data_t* data,
+             const Index_t nrow,
+             Index_t* output_graph,
+             bool return_distances,
+             DistData_t* output_distances);
   ~GNND()    = default;
   using ID_t = InternalID_t<Index_t>;
 
@@ -1212,7 +1216,11 @@ void GNND<Data_t, Index_t>::local_join(cudaStream_t stream)
 }
 
 template <typename Data_t, typename Index_t>
-void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* output_graph, bool return_distances, DistData_t *output_distances)
+void GNND<Data_t, Index_t>::build(Data_t* data,
+                                  const Index_t nrow,
+                                  Index_t* output_graph,
+                                  bool return_distances,
+                                  DistData_t* output_distances)
 {
   using input_t = typename std::remove_const<Data_t>::type;
 
@@ -1339,7 +1347,7 @@ void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* out
   // Reuse graph_.h_dists as the buffer for shrink the lists in graph
   static_assert(sizeof(decltype(*(graph_.h_dists.data_handle()))) >= sizeof(Index_t));
 
-  if(return_distances) {
+  if (return_distances) {
     raft::copy(output_distances,
                graph_.h_dists.data_handle(),
                nrow_ * build_config_.node_degree,
@@ -1412,9 +1420,9 @@ void build(raft::resources const& res,
 
   auto int_graph = raft::make_host_matrix<int, int64_t, row_major>(
     dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
-  
-  auto distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(0,0);
-  if(params.return_distances) {
+
+  auto distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0);
+  if (params.return_distances) {
     distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(
       dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
   }
@@ -1427,16 +1435,21 @@ void build(raft::resources const& res,
                            .termination_threshold = params.termination_threshold};
 
   GNND<const T, int> nnd(res, build_config);
-  nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle(), params.return_distances, distances_graph.data_handle());
+  nnd.build(dataset.data_handle(),
+            dataset.extent(0),
+            int_graph.data_handle(),
+            params.return_distances,
+            distances_graph.data_handle());
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
     for (size_t j = 0; j < graph_degree; j++) {
       auto graph                  = idx.graph().data_handle();
       graph[i * graph_degree + j] = int_graph.data_handle()[i * extended_graph_degree + j];
-      if(params.return_distances) {
+      if (params.return_distances) {
         auto dist_graph = idx.distances().data_handle();
-        dist_graph[i * graph_degree + j] = distances_graph.data_handle()[i * extended_graph_degree + j];
+        dist_graph[i * graph_degree + j] =
+          distances_graph.data_handle()[i * extended_graph_degree + j];
       }
     }
   }
diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index 36ac32eb24..edf0a890ad 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -121,8 +121,14 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
       }
 
       double min_recall = ps.min_recall;
-      EXPECT_TRUE(eval_neighbours(
-        indices_naive, indices_NNDescent, distances_naive, distances_NNDescent, ps.n_rows, ps.graph_degree, 0.001, min_recall));
+      EXPECT_TRUE(eval_neighbours(indices_naive,
+                                  indices_NNDescent,
+                                  distances_naive,
+                                  distances_NNDescent,
+                                  ps.n_rows,
+                                  ps.graph_degree,
+                                  0.001,
+                                  min_recall));
     }
   }
 

From 11379f76598d3622a9cb5a83a1a15093144f1c58 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Fri, 31 May 2024 19:19:35 +0000
Subject: [PATCH 06/25] change return_distances to bool

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 626adcaa90..5d79179d46 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -52,7 +52,7 @@ struct index_params : ann::index_params {
   size_t intermediate_graph_degree = 128;     // Degree of input graph for pruning.
   size_t max_iterations            = 20;      // Number of nn-descent iterations.
   float termination_threshold      = 0.0001;  // Termination threshold of nn-descent.
-  int return_distances             = 0;       // return distances if 1
+  bool return_distances             = false;       // return distances if true
 };
 
 /**

From e993d7c34658b84aa68890540b13297012e7bfec Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Fri, 31 May 2024 22:13:55 +0000
Subject: [PATCH 07/25] change distances to optional

---
 .../raft/neighbors/detail/nn_descent.cuh      |  3 ++-
 .../raft/neighbors/nn_descent_types.hpp       | 27 ++++++++++++++-----
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index c617c05c6a..0694cd8849 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -45,6 +45,7 @@
 #include <omp.h>
 
 #include <limits>
+#include <optional>
 #include <queue>
 #include <random>
 
@@ -1475,7 +1476,7 @@ index<IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  index<IdxT> idx{res, dataset.extent(0), static_cast<int64_t>(graph_degree)};
+  index<IdxT> idx{res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
 
   build(res, params, dataset, idx);
 
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 5d79179d46..08ae310333 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -24,6 +24,7 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/distance/distance_types.hpp>
+#include <optional>
 
 namespace raft::neighbors::experimental::nn_descent {
 using DistData_t = float;
@@ -82,14 +83,18 @@ struct index : ann::index {
    * @param n_rows number of rows in knn-graph
    * @param n_cols number of cols in knn-graph
    */
-  index(raft::resources const& res, int64_t n_rows, int64_t n_cols)
+  index(raft::resources const& res, int64_t n_rows, int64_t n_cols, bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(n_rows, n_cols)},
-      distances_{raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols)},
-      graph_view_{graph_.view()}
+      graph_view_{graph_.view()},
+      return_distances_(return_distances)
   {
+    if(return_distances) {
+      distances_ = raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols);
+      distances_view_ = distances_.value().view();
+    }
   }
 
   /**
@@ -103,13 +108,17 @@ struct index : ann::index {
    * @param graph_view raft::host_matrix_view<IdxT, int64_t, raft::row_major> for storing knn-graph
    */
   index(raft::resources const& res,
-        raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view)
+        raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
+        std::optional<raft::host_matrix_view<DistData_t, int64_t, raft::row_major>> distances_view = std::nullopt,
+        bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
       distances_{raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0)},
-      graph_view_{graph_view}
+      graph_view_{graph_view},
+      distances_view_(distances_view),
+      return_distances_(return_distances)
   {
   }
 
@@ -139,7 +148,9 @@ struct index : ann::index {
 
   [[nodiscard]] inline auto distances() noexcept -> host_matrix_view<DistData_t, int64_t, row_major>
   {
-    return distances_.view();
+    assert(return_distances_);
+    assert(distances_view_.has_value());
+    return distances_view_.value();
   }
 
   // Don't allow copying the index for performance reasons (try avoiding copying data)
@@ -153,9 +164,11 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
-  raft::host_matrix<DistData_t, int64_t, row_major> distances_;
+  std::optional<raft::host_matrix<DistData_t, int64_t, row_major>> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
+  std::optional<raft::host_matrix_view<DistData_t, int64_t, row_major>> distances_view_;
+  bool return_distances_;
 };
 
 /** @} */

From 4578dc0f1a671f9d42238a32829dd2f8d7ac19a8 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Fri, 31 May 2024 22:20:20 +0000
Subject: [PATCH 08/25] fix styling:

---
 cpp/include/raft/neighbors/detail/nn_descent.cuh |  3 ++-
 cpp/include/raft/neighbors/nn_descent_types.hpp  | 10 ++++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index 0694cd8849..dd10cb6111 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1476,7 +1476,8 @@ index<IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  index<IdxT> idx{res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
+  index<IdxT> idx{
+    res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
 
   build(res, params, dataset, idx);
 
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 08ae310333..598f2ef448 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -24,6 +24,7 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/distance/distance_types.hpp>
+
 #include <optional>
 
 namespace raft::neighbors::experimental::nn_descent {
@@ -53,7 +54,7 @@ struct index_params : ann::index_params {
   size_t intermediate_graph_degree = 128;     // Degree of input graph for pruning.
   size_t max_iterations            = 20;      // Number of nn-descent iterations.
   float termination_threshold      = 0.0001;  // Termination threshold of nn-descent.
-  bool return_distances             = false;       // return distances if true
+  bool return_distances            = false;   // return distances if true
 };
 
 /**
@@ -91,8 +92,8 @@ struct index : ann::index {
       graph_view_{graph_.view()},
       return_distances_(return_distances)
   {
-    if(return_distances) {
-      distances_ = raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols);
+    if (return_distances) {
+      distances_      = raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols);
       distances_view_ = distances_.value().view();
     }
   }
@@ -109,7 +110,8 @@ struct index : ann::index {
    */
   index(raft::resources const& res,
         raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-        std::optional<raft::host_matrix_view<DistData_t, int64_t, raft::row_major>> distances_view = std::nullopt,
+        std::optional<raft::host_matrix_view<DistData_t, int64_t, raft::row_major>> distances_view =
+          std::nullopt,
         bool return_distances = false)
     : ann::index(),
       res_{res},

From c9048a5f9c73d5650df650b510f114c367338191 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 5 Jun 2024 19:21:12 +0000
Subject: [PATCH 09/25] handle bad access error

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 598f2ef448..01e632cd34 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -150,9 +150,11 @@ struct index : ann::index {
 
   [[nodiscard]] inline auto distances() noexcept -> host_matrix_view<DistData_t, int64_t, row_major>
   {
-    assert(return_distances_);
-    assert(distances_view_.has_value());
-    return distances_view_.value();
+    if (distances_view_.has_value()) {
+      return distances_view_.value();
+    } else {
+      return raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0).view();
+    }
   }
 
   // Don't allow copying the index for performance reasons (try avoiding copying data)

From bdb74a5df84ab38c9c58c496c90c0aea6528d878 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 5 Jun 2024 21:21:48 +0000
Subject: [PATCH 10/25] remove unnecessary dist allocation

---
 .../raft/neighbors/detail/nn_descent.cuh      | 25 +++++++------------
 .../raft/neighbors/nn_descent_types.hpp       |  3 +++
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index dd10cb6111..333c12d303 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -218,6 +218,7 @@ struct BuildConfig {
   // If internal_node_degree == 0, the value of node_degree will be assigned to it
   size_t max_iterations{50};
   float termination_threshold{0.0001};
+  size_t output_graph_degree{32};
 };
 
 template <typename Index_t>
@@ -1349,10 +1350,12 @@ void GNND<Data_t, Index_t>::build(Data_t* data,
   static_assert(sizeof(decltype(*(graph_.h_dists.data_handle()))) >= sizeof(Index_t));
 
   if (return_distances) {
-    raft::copy(output_distances,
-               graph_.h_dists.data_handle(),
-               nrow_ * build_config_.node_degree,
+    for (size_t i = 0; i < (size_t)nrow_; i++) {
+      raft::copy(output_distances + i * build_config_.output_graph_degree,
+               graph_.h_dists.data_handle() + i * build_config_.node_degree,
+               build_config_.output_graph_degree,
                raft::resource::get_cuda_stream(res));
+    }
   }
 
   Index_t* graph_shrink_buffer = (Index_t*)graph_.h_dists.data_handle();
@@ -1422,36 +1425,26 @@ void build(raft::resources const& res,
   auto int_graph = raft::make_host_matrix<int, int64_t, row_major>(
     dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
 
-  auto distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0);
-  if (params.return_distances) {
-    distances_graph = raft::make_host_matrix<DistData_t, int64_t, row_major>(
-      dataset.extent(0), static_cast<int64_t>(extended_graph_degree));
-  }
-
   BuildConfig build_config{.max_dataset_size      = static_cast<size_t>(dataset.extent(0)),
                            .dataset_dim           = static_cast<size_t>(dataset.extent(1)),
                            .node_degree           = extended_graph_degree,
                            .internal_node_degree  = extended_intermediate_degree,
                            .max_iterations        = params.max_iterations,
-                           .termination_threshold = params.termination_threshold};
+                           .termination_threshold = params.termination_threshold,
+                           .output_graph_degree    = params.graph_degree};
 
   GNND<const T, int> nnd(res, build_config);
   nnd.build(dataset.data_handle(),
             dataset.extent(0),
             int_graph.data_handle(),
             params.return_distances,
-            distances_graph.data_handle());
+            idx.distances().data_handle());
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
     for (size_t j = 0; j < graph_degree; j++) {
       auto graph                  = idx.graph().data_handle();
       graph[i * graph_degree + j] = int_graph.data_handle()[i * extended_graph_degree + j];
-      if (params.return_distances) {
-        auto dist_graph = idx.distances().data_handle();
-        dist_graph[i * graph_degree + j] =
-          distances_graph.data_handle()[i * extended_graph_degree + j];
-      }
     }
   }
 }
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 01e632cd34..9e35e1d65c 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -122,6 +122,9 @@ struct index : ann::index {
       distances_view_(distances_view),
       return_distances_(return_distances)
   {
+    if(!distances_view.has_value()) {
+      distances_view_ = distances_.value().view();
+    }
   }
 
   /** Distance metric used for clustering. */

From c71e97a646bc6916a83cfe0fd25a6d2fee215db8 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 5 Jun 2024 22:28:55 +0000
Subject: [PATCH 11/25] change to device matrix

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 9e35e1d65c..efc773e275 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -20,6 +20,8 @@
 
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/host_mdspan.hpp>
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_mdspan.hpp>
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
@@ -93,7 +95,7 @@ struct index : ann::index {
       return_distances_(return_distances)
   {
     if (return_distances) {
-      distances_      = raft::make_host_matrix<DistData_t, int64_t, row_major>(n_rows, n_cols);
+      distances_      = raft::make_device_matrix<DistData_t, int64_t>(res_, n_rows, n_cols);
       distances_view_ = distances_.value().view();
     }
   }
@@ -110,14 +112,14 @@ struct index : ann::index {
    */
   index(raft::resources const& res,
         raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-        std::optional<raft::host_matrix_view<DistData_t, int64_t, raft::row_major>> distances_view =
+        std::optional<raft::device_matrix_view<DistData_t, int64_t, row_major>> distances_view =
           std::nullopt,
         bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
-      distances_{raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0)},
+      distances_{raft::make_device_matrix<DistData_t, int64_t>(res_, 0, 0)},
       graph_view_{graph_view},
       distances_view_(distances_view),
       return_distances_(return_distances)
@@ -151,12 +153,12 @@ struct index : ann::index {
     return graph_view_;
   }
 
-  [[nodiscard]] inline auto distances() noexcept -> host_matrix_view<DistData_t, int64_t, row_major>
+  [[nodiscard]] inline auto distances() noexcept -> device_matrix_view<DistData_t, int64_t, row_major>
   {
     if (distances_view_.has_value()) {
       return distances_view_.value();
     } else {
-      return raft::make_host_matrix<DistData_t, int64_t, row_major>(0, 0).view();
+      return raft::make_device_matrix<DistData_t, int64_t>(res_, 0, 0).view();
     }
   }
 
@@ -171,10 +173,10 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
-  std::optional<raft::host_matrix<DistData_t, int64_t, row_major>> distances_;
+  std::optional<raft::device_matrix<DistData_t, int64_t, row_major>> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
-  std::optional<raft::host_matrix_view<DistData_t, int64_t, row_major>> distances_view_;
+  std::optional<raft::device_matrix_view<DistData_t, int64_t, row_major>> distances_view_;
   bool return_distances_;
 };
 

From 0327ce591a246499676427d03abbc9b42528b5a1 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 5 Jun 2024 23:44:35 +0000
Subject: [PATCH 12/25] change template param for index

---
 .../neighbors/detail/cagra/cagra_build.cuh     |  2 +-
 .../raft/neighbors/detail/nn_descent.cuh       |  6 +++---
 cpp/include/raft/neighbors/nn_descent.cuh      |  8 ++++----
 .../raft/neighbors/nn_descent_types.hpp        | 18 +++++++++---------
 4 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
index 40dcf68e68..c558f90c84 100644
--- a/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
@@ -240,7 +240,7 @@ void build_knn_graph(raft::resources const& res,
                      raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph,
                      experimental::nn_descent::index_params build_params)
 {
-  auto nn_descent_idx = experimental::nn_descent::index<IdxT>(res, knn_graph);
+  auto nn_descent_idx = experimental::nn_descent::index<float, IdxT>(res, knn_graph);
   experimental::nn_descent::build<DataT, IdxT>(res, build_params, dataset, nn_descent_idx);
 
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index 333c12d303..a7aa956c51 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1391,7 +1391,7 @@ template <typename T,
 void build(raft::resources const& res,
            const index_params& params,
            mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset,
-           index<IdxT>& idx)
+           index<DistData_t, IdxT>& idx)
 {
   RAFT_EXPECTS(dataset.extent(0) < std::numeric_limits<int>::max() - 1,
                "The dataset size for GNND should be less than %d",
@@ -1453,7 +1453,7 @@ template <typename T,
           typename IdxT = uint32_t,
           typename Accessor =
             host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
-index<IdxT> build(raft::resources const& res,
+index<DistData_t, IdxT> build(raft::resources const& res,
                   const index_params& params,
                   mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
 {
@@ -1469,7 +1469,7 @@ index<IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  index<IdxT> idx{
+  index<DistData_t, IdxT> idx{
     res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
 
   build(res, params, dataset, idx);
diff --git a/cpp/include/raft/neighbors/nn_descent.cuh b/cpp/include/raft/neighbors/nn_descent.cuh
index ceb5ae5643..62c03e2b19 100644
--- a/cpp/include/raft/neighbors/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/nn_descent.cuh
@@ -56,7 +56,7 @@ namespace raft::neighbors::experimental::nn_descent {
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<IdxT> build(raft::resources const& res,
+index<detail::DistData_t, IdxT> build(raft::resources const& res,
                   index_params const& params,
                   raft::device_matrix_view<const T, int64_t, row_major> dataset)
 {
@@ -97,7 +97,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::device_matrix_view<const T, int64_t, row_major> dataset,
-           index<IdxT>& idx)
+           index<detail::DistData_t, IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }
@@ -130,7 +130,7 @@ void build(raft::resources const& res,
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<IdxT> build(raft::resources const& res,
+index<detail::DistData_t, IdxT> build(raft::resources const& res,
                   index_params const& params,
                   raft::host_matrix_view<const T, int64_t, row_major> dataset)
 {
@@ -171,7 +171,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::host_matrix_view<const T, int64_t, row_major> dataset,
-           index<IdxT>& idx)
+           index<detail::DistData_t, IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index efc773e275..b353d7fd02 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -30,7 +30,7 @@
 #include <optional>
 
 namespace raft::neighbors::experimental::nn_descent {
-using DistData_t = float;
+// using DistData_t = float;
 /**
  * @ingroup nn-descent
  * @{
@@ -72,7 +72,7 @@ struct index_params : ann::index_params {
  *
  * @tparam IdxT dtype to be used for constructing knn-graph
  */
-template <typename IdxT>
+template <typename T, typename IdxT>
 struct index : ann::index {
  public:
   /**
@@ -95,7 +95,7 @@ struct index : ann::index {
       return_distances_(return_distances)
   {
     if (return_distances) {
-      distances_      = raft::make_device_matrix<DistData_t, int64_t>(res_, n_rows, n_cols);
+      distances_      = raft::make_device_matrix<T, int64_t>(res_, n_rows, n_cols);
       distances_view_ = distances_.value().view();
     }
   }
@@ -112,14 +112,14 @@ struct index : ann::index {
    */
   index(raft::resources const& res,
         raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-        std::optional<raft::device_matrix_view<DistData_t, int64_t, row_major>> distances_view =
+        std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view =
           std::nullopt,
         bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
-      distances_{raft::make_device_matrix<DistData_t, int64_t>(res_, 0, 0)},
+      distances_{raft::make_device_matrix<T, int64_t>(res_, 0, 0)},
       graph_view_{graph_view},
       distances_view_(distances_view),
       return_distances_(return_distances)
@@ -153,12 +153,12 @@ struct index : ann::index {
     return graph_view_;
   }
 
-  [[nodiscard]] inline auto distances() noexcept -> device_matrix_view<DistData_t, int64_t, row_major>
+  [[nodiscard]] inline auto distances() noexcept -> device_matrix_view<T, int64_t, row_major>
   {
     if (distances_view_.has_value()) {
       return distances_view_.value();
     } else {
-      return raft::make_device_matrix<DistData_t, int64_t>(res_, 0, 0).view();
+      return raft::make_device_matrix<T, int64_t>(res_, 0, 0).view();
     }
   }
 
@@ -173,10 +173,10 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
-  std::optional<raft::device_matrix<DistData_t, int64_t, row_major>> distances_;
+  std::optional<raft::device_matrix<T, int64_t, row_major>> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
-  std::optional<raft::device_matrix_view<DistData_t, int64_t, row_major>> distances_view_;
+  std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view_;
   bool return_distances_;
 };
 

From 3f49752e2937a0d7a4c98ddc7a9b95c6417245b1 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 6 Jun 2024 01:55:22 +0000
Subject: [PATCH 13/25] update test

---
 cpp/test/neighbors/ann_nn_descent.cuh | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index edf0a890ad..dc8060c4d4 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -17,6 +17,7 @@
 
 #include "../test_utils.cuh"
 #include "ann_utils.cuh"
+#include "raft/util/cudart_utils.hpp"
 
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/neighbors/nn_descent.cuh>
@@ -94,6 +95,7 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
         index_params.graph_degree              = ps.graph_degree;
         index_params.intermediate_graph_degree = 2 * ps.graph_degree;
         index_params.max_iterations            = 100;
+        index_params.return_distances          = true;
 
         auto database_view = raft::make_device_matrix_view<const DataT, int64_t>(
           (const DataT*)database.data(), ps.n_rows, ps.dim);
@@ -105,16 +107,12 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
             auto database_host_view = raft::make_host_matrix_view<const DataT, int64_t>(
               (const DataT*)database_host.data_handle(), ps.n_rows, ps.dim);
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_host_view);
-            update_host(
-              indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            update_host(
-              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            raft::copy(distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           } else {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_view);
-            update_host(
-              indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            update_host(
-              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            raft::copy(distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           };
         }
         resource::sync_stream(handle_);

From a6a8ad2e367e5280b302a256ff86be30425015d6 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 6 Jun 2024 02:22:22 +0000
Subject: [PATCH 14/25] remove comment

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index b353d7fd02..93aca6e2a1 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -30,7 +30,6 @@
 #include <optional>
 
 namespace raft::neighbors::experimental::nn_descent {
-// using DistData_t = float;
 /**
  * @ingroup nn-descent
  * @{

From 2cfda3c982540c86683717989319884e6d63204e Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 6 Jun 2024 02:24:22 +0000
Subject: [PATCH 15/25] fix styling

---
 .../raft/neighbors/detail/nn_descent.cuh       | 12 ++++++------
 cpp/include/raft/neighbors/nn_descent.cuh      | 10 +++++-----
 .../raft/neighbors/nn_descent_types.hpp        | 18 ++++++++----------
 cpp/test/neighbors/ann_nn_descent.cuh          | 12 ++++++++----
 4 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index a7aa956c51..53419f6db6 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1352,9 +1352,9 @@ void GNND<Data_t, Index_t>::build(Data_t* data,
   if (return_distances) {
     for (size_t i = 0; i < (size_t)nrow_; i++) {
       raft::copy(output_distances + i * build_config_.output_graph_degree,
-               graph_.h_dists.data_handle() + i * build_config_.node_degree,
-               build_config_.output_graph_degree,
-               raft::resource::get_cuda_stream(res));
+                 graph_.h_dists.data_handle() + i * build_config_.node_degree,
+                 build_config_.output_graph_degree,
+                 raft::resource::get_cuda_stream(res));
     }
   }
 
@@ -1431,7 +1431,7 @@ void build(raft::resources const& res,
                            .internal_node_degree  = extended_intermediate_degree,
                            .max_iterations        = params.max_iterations,
                            .termination_threshold = params.termination_threshold,
-                           .output_graph_degree    = params.graph_degree};
+                           .output_graph_degree   = params.graph_degree};
 
   GNND<const T, int> nnd(res, build_config);
   nnd.build(dataset.data_handle(),
@@ -1454,8 +1454,8 @@ template <typename T,
           typename Accessor =
             host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
 index<DistData_t, IdxT> build(raft::resources const& res,
-                  const index_params& params,
-                  mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
+                              const index_params& params,
+                              mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
diff --git a/cpp/include/raft/neighbors/nn_descent.cuh b/cpp/include/raft/neighbors/nn_descent.cuh
index 62c03e2b19..a992836954 100644
--- a/cpp/include/raft/neighbors/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/nn_descent.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-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.
@@ -57,8 +57,8 @@ namespace raft::neighbors::experimental::nn_descent {
  */
 template <typename T, typename IdxT = uint32_t>
 index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                  index_params const& params,
-                  raft::device_matrix_view<const T, int64_t, row_major> dataset)
+                                      index_params const& params,
+                                      raft::device_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
@@ -131,8 +131,8 @@ void build(raft::resources const& res,
  */
 template <typename T, typename IdxT = uint32_t>
 index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                  index_params const& params,
-                  raft::host_matrix_view<const T, int64_t, row_major> dataset)
+                                      index_params const& params,
+                                      raft::host_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 93aca6e2a1..a9a83b25e1 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -18,10 +18,10 @@
 
 #include "ann_types.hpp"
 
-#include <raft/core/host_mdarray.hpp>
-#include <raft/core/host_mdspan.hpp>
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_mdspan.hpp>
+#include <raft/core/host_mdarray.hpp>
+#include <raft/core/host_mdspan.hpp>
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
@@ -109,11 +109,11 @@ struct index : ann::index {
    * @param res raft::resources is an object mangaging resources
    * @param graph_view raft::host_matrix_view<IdxT, int64_t, raft::row_major> for storing knn-graph
    */
-  index(raft::resources const& res,
-        raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-        std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view =
-          std::nullopt,
-        bool return_distances = false)
+  index(
+    raft::resources const& res,
+    raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
+    std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view = std::nullopt,
+    bool return_distances                                                         = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
@@ -123,9 +123,7 @@ struct index : ann::index {
       distances_view_(distances_view),
       return_distances_(return_distances)
   {
-    if(!distances_view.has_value()) {
-      distances_view_ = distances_.value().view();
-    }
+    if (!distances_view.has_value()) { distances_view_ = distances_.value().view(); }
   }
 
   /** Distance metric used for clustering. */
diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index dc8060c4d4..fe3e28b9b8 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -107,12 +107,16 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
             auto database_host_view = raft::make_host_matrix_view<const DataT, int64_t>(
               (const DataT*)database_host.data_handle(), ps.n_rows, ps.dim);
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_host_view);
-            raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            raft::copy(distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            raft::copy(
+              indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            raft::copy(
+              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           } else {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_view);
-            raft::copy(indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            raft::copy(distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            raft::copy(
+              indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
+            raft::copy(
+              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
           };
         }
         resource::sync_stream(handle_);

From 2efae1a96fd50b0f5535ee300fd6fb1defa52ae5 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 6 Jun 2024 17:16:42 +0000
Subject: [PATCH 16/25] fix header

---
 cpp/test/neighbors/ann_nn_descent.cuh | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index fe3e28b9b8..12b50ad532 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -15,12 +15,11 @@
  */
 #pragma once
 
-#include "../test_utils.cuh"
 #include "ann_utils.cuh"
-#include "raft/util/cudart_utils.hpp"
 
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/neighbors/nn_descent.cuh>
+#include <raft/util/cudart_utils.hpp>
 #include <raft/util/itertools.hpp>
 
 #include <raft_internal/neighbors/naive_knn.cuh>

From 931158c534118b7a42d23902e8ab1e1f9cde2452 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 6 Jun 2024 20:43:53 +0000
Subject: [PATCH 17/25] add documentation for return_distances

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index a9a83b25e1..042a7fc82c 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -84,6 +84,7 @@ struct index : ann::index {
    * @param res raft::resources is an object mangaging resources
    * @param n_rows number of rows in knn-graph
    * @param n_cols number of cols in knn-graph
+   * @param return_distances whether to allocate and get distances information
    */
   index(raft::resources const& res, int64_t n_rows, int64_t n_cols, bool return_distances = false)
     : ann::index(),

From 9f17b5c1287258034866588ef0197a12a53eeb85 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Fri, 7 Jun 2024 19:29:41 +0000
Subject: [PATCH 18/25] add documentation

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 042a7fc82c..94a6f3db4b 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -109,6 +109,9 @@ struct index : ann::index {
    *
    * @param res raft::resources is an object mangaging resources
    * @param graph_view raft::host_matrix_view<IdxT, int64_t, raft::row_major> for storing knn-graph
+   * @param distances_view std::optional<raft::device_matrix_view<T, int64_t, row_major>> for
+   * storing knn-graph distances
+   * @param return_distances whether to allocate and get distances information
    */
   index(
     raft::resources const& res,
@@ -151,6 +154,7 @@ struct index : ann::index {
     return graph_view_;
   }
 
+  /** neighborhood graph distances [size, graph-degree] */
   [[nodiscard]] inline auto distances() noexcept -> device_matrix_view<T, int64_t, row_major>
   {
     if (distances_view_.has_value()) {

From e28c2f9b03a12022c491f4bb9e7db3518bb9209a Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Mon, 10 Jun 2024 16:19:49 +0000
Subject: [PATCH 19/25] add tparam doc

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 94a6f3db4b..6fe859721f 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -69,6 +69,7 @@ struct index_params : ann::index_params {
  * International Conference on Information & Knowledge Management (CIKM '21). Association for
  * Computing Machinery, New York, NY, USA, 1929–1938. https://doi.org/10.1145/3459637.3482344
  *
+ * @tparam T dtype to be used for constructing distances graph
  * @tparam IdxT dtype to be used for constructing knn-graph
  */
 template <typename T, typename IdxT>

From fd6544262272fee98985515421e8917496664dab Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Wed, 12 Jun 2024 21:35:55 +0000
Subject: [PATCH 20/25] remove redundancy

---
 cpp/include/raft/neighbors/nn_descent_types.hpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 6fe859721f..44e99cd355 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -128,7 +128,6 @@ struct index : ann::index {
       distances_view_(distances_view),
       return_distances_(return_distances)
   {
-    if (!distances_view.has_value()) { distances_view_ = distances_.value().view(); }
   }
 
   /** Distance metric used for clustering. */

From 3ce8cdf7414c5342157f1786ae3be9c9aca20481 Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 13 Jun 2024 00:04:23 +0000
Subject: [PATCH 21/25] return optional for distances()

---
 cpp/include/raft/neighbors/detail/nn_descent.cuh | 10 ++++++++--
 cpp/include/raft/neighbors/nn_descent_types.hpp  |  8 ++------
 cpp/test/neighbors/ann_nn_descent.cuh            | 13 +++++++++----
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index 53419f6db6..d53765c373 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1434,11 +1434,17 @@ void build(raft::resources const& res,
                            .output_graph_degree   = params.graph_degree};
 
   GNND<const T, int> nnd(res, build_config);
-  nnd.build(dataset.data_handle(),
+
+  if(idx.distances().has_value()) {
+    nnd.build(dataset.data_handle(),
             dataset.extent(0),
             int_graph.data_handle(),
             params.return_distances,
-            idx.distances().data_handle());
+            idx.distances().value().data_handle());
+  } else {
+    RAFT_EXPECTS(false, "Distance view not allocated. Using NN Descent requires return_distances to be true");
+  }
+  
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 44e99cd355..4bad85908d 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -155,13 +155,9 @@ struct index : ann::index {
   }
 
   /** neighborhood graph distances [size, graph-degree] */
-  [[nodiscard]] inline auto distances() noexcept -> device_matrix_view<T, int64_t, row_major>
+  [[nodiscard]] inline auto distances() noexcept -> std::optional<device_matrix_view<T, int64_t, row_major>>
   {
-    if (distances_view_.has_value()) {
-      return distances_view_.value();
-    } else {
-      return raft::make_device_matrix<T, int64_t>(res_, 0, 0).view();
-    }
+    return distances_view_;
   }
 
   // Don't allow copying the index for performance reasons (try avoiding copying data)
diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index 12b50ad532..953da2d42b 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -108,14 +108,19 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_host_view);
             raft::copy(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            raft::copy(
-              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            if(index.distances().has_value()) {
+              raft::copy(
+              distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_);
+            }
+            
           } else {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_view);
             raft::copy(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            raft::copy(
-              distances_NNDescent.data(), index.distances().data_handle(), queries_size, stream_);
+            if(index.distances().has_value()) {
+              raft::copy(
+              distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_);
+            }
           };
         }
         resource::sync_stream(handle_);

From 4ec3ab0a8583abc2c8ed98f0cc9ff36a3278ed0e Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 13 Jun 2024 00:22:46 +0000
Subject: [PATCH 22/25] fix styling

---
 .../raft/neighbors/detail/nn_descent.cuh       | 14 +++++++-------
 .../raft/neighbors/nn_descent_types.hpp        |  3 ++-
 cpp/test/neighbors/ann_nn_descent.cuh          | 18 +++++++++++-------
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index d53765c373..a2dc9c00d7 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1435,16 +1435,16 @@ void build(raft::resources const& res,
 
   GNND<const T, int> nnd(res, build_config);
 
-  if(idx.distances().has_value()) {
+  if (idx.distances().has_value()) {
     nnd.build(dataset.data_handle(),
-            dataset.extent(0),
-            int_graph.data_handle(),
-            params.return_distances,
-            idx.distances().value().data_handle());
+              dataset.extent(0),
+              int_graph.data_handle(),
+              params.return_distances,
+              idx.distances().value().data_handle());
   } else {
-    RAFT_EXPECTS(false, "Distance view not allocated. Using NN Descent requires return_distances to be true");
+    RAFT_EXPECTS(
+      false, "Distance view not allocated. Using NN Descent requires return_distances to be true");
   }
-  
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 4bad85908d..9838047bb5 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -155,7 +155,8 @@ struct index : ann::index {
   }
 
   /** neighborhood graph distances [size, graph-degree] */
-  [[nodiscard]] inline auto distances() noexcept -> std::optional<device_matrix_view<T, int64_t, row_major>>
+  [[nodiscard]] inline auto distances() noexcept
+    -> std::optional<device_matrix_view<T, int64_t, row_major>>
   {
     return distances_view_;
   }
diff --git a/cpp/test/neighbors/ann_nn_descent.cuh b/cpp/test/neighbors/ann_nn_descent.cuh
index 953da2d42b..f74cadb415 100644
--- a/cpp/test/neighbors/ann_nn_descent.cuh
+++ b/cpp/test/neighbors/ann_nn_descent.cuh
@@ -108,18 +108,22 @@ class AnnNNDescentTest : public ::testing::TestWithParam<AnnNNDescentInputs> {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_host_view);
             raft::copy(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            if(index.distances().has_value()) {
-              raft::copy(
-              distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_);
+            if (index.distances().has_value()) {
+              raft::copy(distances_NNDescent.data(),
+                         index.distances().value().data_handle(),
+                         queries_size,
+                         stream_);
             }
-            
+
           } else {
             auto index = nn_descent::build<DataT, IdxT>(handle_, index_params, database_view);
             raft::copy(
               indices_NNDescent.data(), index.graph().data_handle(), queries_size, stream_);
-            if(index.distances().has_value()) {
-              raft::copy(
-              distances_NNDescent.data(), index.distances().value().data_handle(), queries_size, stream_);
+            if (index.distances().has_value()) {
+              raft::copy(distances_NNDescent.data(),
+                         index.distances().value().data_handle(),
+                         queries_size,
+                         stream_);
             }
           };
         }

From b0347a7a5a77af55d2aeb7484edb299a2332765a Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 13 Jun 2024 15:16:19 +0000
Subject: [PATCH 23/25] remove raft_expects

---
 cpp/include/raft/neighbors/detail/nn_descent.cuh | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index a2dc9c00d7..8d156dfe59 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1435,16 +1435,12 @@ void build(raft::resources const& res,
 
   GNND<const T, int> nnd(res, build_config);
 
-  if (idx.distances().has_value()) {
-    nnd.build(dataset.data_handle(),
-              dataset.extent(0),
-              int_graph.data_handle(),
-              params.return_distances,
-              idx.distances().value().data_handle());
-  } else {
-    RAFT_EXPECTS(
-      false, "Distance view not allocated. Using NN Descent requires return_distances to be true");
-  }
+  // auto distances_graph = raft::make_device_matrix<float, int64_t>(res, 0, 0)
+  nnd.build(dataset.data_handle(),
+            dataset.extent(0),
+            int_graph.data_handle(),
+            params.return_distances,
+            idx.distances().value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view()).data_handle());
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {

From 84abd964aa2c3dea25bcb64b6dee36618cbe684a Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 13 Jun 2024 20:32:04 +0000
Subject: [PATCH 24/25] fix type template and bring back raft_expects

---
 .../neighbors/detail/cagra/cagra_build.cuh    |  2 +-
 .../raft/neighbors/detail/nn_descent.cuh      | 29 ++++++++++++-------
 cpp/include/raft/neighbors/nn_descent.cuh     | 18 ++++++------
 .../raft/neighbors/nn_descent_types.hpp       | 23 +++++++--------
 4 files changed, 39 insertions(+), 33 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
index c558f90c84..40dcf68e68 100644
--- a/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
@@ -240,7 +240,7 @@ void build_knn_graph(raft::resources const& res,
                      raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph,
                      experimental::nn_descent::index_params build_params)
 {
-  auto nn_descent_idx = experimental::nn_descent::index<float, IdxT>(res, knn_graph);
+  auto nn_descent_idx = experimental::nn_descent::index<IdxT>(res, knn_graph);
   experimental::nn_descent::build<DataT, IdxT>(res, build_params, dataset, nn_descent_idx);
 
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index 8d156dfe59..8fbf09e7fe 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1391,7 +1391,7 @@ template <typename T,
 void build(raft::resources const& res,
            const index_params& params,
            mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset,
-           index<DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   RAFT_EXPECTS(dataset.extent(0) < std::numeric_limits<int>::max() - 1,
                "The dataset size for GNND should be less than %d",
@@ -1435,12 +1435,19 @@ void build(raft::resources const& res,
 
   GNND<const T, int> nnd(res, build_config);
 
-  // auto distances_graph = raft::make_device_matrix<float, int64_t>(res, 0, 0)
-  nnd.build(dataset.data_handle(),
-            dataset.extent(0),
-            int_graph.data_handle(),
-            params.return_distances,
-            idx.distances().value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view()).data_handle());
+  if (idx.distances().has_value() || !params.return_distances) {
+    nnd.build(dataset.data_handle(),
+              dataset.extent(0),
+              int_graph.data_handle(),
+              params.return_distances,
+              idx.distances()
+                .value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view())
+                .data_handle());
+  } else {
+    RAFT_EXPECTS(false,
+                 "Distance view not allocated. Using return_distances set to true requires "
+                 "distance view to be allocated.");
+  }
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
@@ -1455,9 +1462,9 @@ template <typename T,
           typename IdxT = uint32_t,
           typename Accessor =
             host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
-index<DistData_t, IdxT> build(raft::resources const& res,
-                              const index_params& params,
-                              mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
+index<IdxT> build(raft::resources const& res,
+                  const index_params& params,
+                  mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
@@ -1471,7 +1478,7 @@ index<DistData_t, IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  index<DistData_t, IdxT> idx{
+  index<IdxT> idx{
     res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
 
   build(res, params, dataset, idx);
diff --git a/cpp/include/raft/neighbors/nn_descent.cuh b/cpp/include/raft/neighbors/nn_descent.cuh
index a992836954..ceb5ae5643 100644
--- a/cpp/include/raft/neighbors/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/nn_descent.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -56,9 +56,9 @@ namespace raft::neighbors::experimental::nn_descent {
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                                      index_params const& params,
-                                      raft::device_matrix_view<const T, int64_t, row_major> dataset)
+index<IdxT> build(raft::resources const& res,
+                  index_params const& params,
+                  raft::device_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
@@ -97,7 +97,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::device_matrix_view<const T, int64_t, row_major> dataset,
-           index<detail::DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }
@@ -130,9 +130,9 @@ void build(raft::resources const& res,
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                                      index_params const& params,
-                                      raft::host_matrix_view<const T, int64_t, row_major> dataset)
+index<IdxT> build(raft::resources const& res,
+                  index_params const& params,
+                  raft::host_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
@@ -171,7 +171,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::host_matrix_view<const T, int64_t, row_major> dataset,
-           index<detail::DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }
diff --git a/cpp/include/raft/neighbors/nn_descent_types.hpp b/cpp/include/raft/neighbors/nn_descent_types.hpp
index 9838047bb5..5d23ff2c2e 100644
--- a/cpp/include/raft/neighbors/nn_descent_types.hpp
+++ b/cpp/include/raft/neighbors/nn_descent_types.hpp
@@ -69,10 +69,9 @@ struct index_params : ann::index_params {
  * International Conference on Information & Knowledge Management (CIKM '21). Association for
  * Computing Machinery, New York, NY, USA, 1929–1938. https://doi.org/10.1145/3459637.3482344
  *
- * @tparam T dtype to be used for constructing distances graph
  * @tparam IdxT dtype to be used for constructing knn-graph
  */
-template <typename T, typename IdxT>
+template <typename IdxT>
 struct index : ann::index {
  public:
   /**
@@ -96,7 +95,7 @@ struct index : ann::index {
       return_distances_(return_distances)
   {
     if (return_distances) {
-      distances_      = raft::make_device_matrix<T, int64_t>(res_, n_rows, n_cols);
+      distances_      = raft::make_device_matrix<float, int64_t>(res_, n_rows, n_cols);
       distances_view_ = distances_.value().view();
     }
   }
@@ -114,16 +113,16 @@ struct index : ann::index {
    * storing knn-graph distances
    * @param return_distances whether to allocate and get distances information
    */
-  index(
-    raft::resources const& res,
-    raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-    std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view = std::nullopt,
-    bool return_distances                                                         = false)
+  index(raft::resources const& res,
+        raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
+        std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view =
+          std::nullopt,
+        bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
-      distances_{raft::make_device_matrix<T, int64_t>(res_, 0, 0)},
+      distances_{raft::make_device_matrix<float, int64_t>(res_, 0, 0)},
       graph_view_{graph_view},
       distances_view_(distances_view),
       return_distances_(return_distances)
@@ -156,7 +155,7 @@ struct index : ann::index {
 
   /** neighborhood graph distances [size, graph-degree] */
   [[nodiscard]] inline auto distances() noexcept
-    -> std::optional<device_matrix_view<T, int64_t, row_major>>
+    -> std::optional<device_matrix_view<float, int64_t, row_major>>
   {
     return distances_view_;
   }
@@ -172,10 +171,10 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
-  std::optional<raft::device_matrix<T, int64_t, row_major>> distances_;
+  std::optional<raft::device_matrix<float, int64_t, row_major>> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
-  std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view_;
+  std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view_;
   bool return_distances_;
 };
 

From ec44b4d78dba1f79abb87735b587dfc11424567f Mon Sep 17 00:00:00 2001
From: jinsolp <soleeep99@gmail.com>
Date: Thu, 13 Jun 2024 20:50:40 +0000
Subject: [PATCH 25/25] fix raft_expects

---
 cpp/include/raft/neighbors/detail/nn_descent.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/detail/nn_descent.cuh b/cpp/include/raft/neighbors/detail/nn_descent.cuh
index 8fbf09e7fe..cdfb9d9931 100644
--- a/cpp/include/raft/neighbors/detail/nn_descent.cuh
+++ b/cpp/include/raft/neighbors/detail/nn_descent.cuh
@@ -1444,7 +1444,7 @@ void build(raft::resources const& res,
                 .value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view())
                 .data_handle());
   } else {
-    RAFT_EXPECTS(false,
+    RAFT_EXPECTS(!params.return_distances,
                  "Distance view not allocated. Using return_distances set to true requires "
                  "distance view to be allocated.");
   }