From 2d8f4863fb1b20296bab303d58199c7eb7530edb Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 1 Feb 2022 18:31:12 -0500
Subject: [PATCH 1/5] cuml updates for lap, label, cluster, and spectral apis

---
 cpp/src/hdbscan/condensed_hierarchy.cu    | 2 +-
 cpp/src/hdbscan/detail/extract.cuh        | 3 +--
 cpp/src/hdbscan/detail/membership.cuh     | 2 +-
 cpp/src/hdbscan/detail/reachability.cuh   | 2 --
 cpp/src/hdbscan/detail/select.cuh         | 2 +-
 cpp/src/hdbscan/detail/stabilities.cuh    | 2 +-
 cpp/src/hdbscan/detail/utils.h            | 4 +---
 cpp/src/hierarchy/pw_dist_graph.cuh       | 1 -
 cpp/src/holtwinters/internal/hw_eval.cuh  | 1 -
 cpp/src/holtwinters/internal/hw_optim.cuh | 1 -
 cpp/src/knn/knn.cu                        | 2 +-
 cpp/src/svm/svc_impl.cuh                  | 2 +-
 cpp/src/tsne/fft_tsne.cuh                 | 1 -
 cpp/src_prims/functions/hinge.cuh         | 1 -
 cpp/src_prims/functions/linearReg.cuh     | 1 -
 cpp/src_prims/label/classlabels.cuh       | 2 +-
 cpp/src_prims/linalg/lstsq.cuh            | 1 -
 cpp/src_prims/linalg/rsvd.cuh             | 1 -
 cpp/src_prims/random/make_regression.cuh  | 1 -
 cpp/test/prims/knn_classify.cu            | 2 +-
 20 files changed, 10 insertions(+), 24 deletions(-)

diff --git a/cpp/src/hdbscan/condensed_hierarchy.cu b/cpp/src/hdbscan/condensed_hierarchy.cu
index 64d01c5883..a7a047263f 100644
--- a/cpp/src/hdbscan/condensed_hierarchy.cu
+++ b/cpp/src/hdbscan/condensed_hierarchy.cu
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <cub/cub.cuh>
 
diff --git a/cpp/src/hdbscan/detail/extract.cuh b/cpp/src/hdbscan/detail/extract.cuh
index 4a9429e98b..b5352b731d 100644
--- a/cpp/src/hdbscan/detail/extract.cuh
+++ b/cpp/src/hdbscan/detail/extract.cuh
@@ -21,12 +21,11 @@
 #include "stabilities.cuh"
 #include "utils.h"
 
-#include <label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 
 #include <raft/cudart_utils.h>
-#include <raft/label/classlabels.cuh>
 #include <raft/sparse/convert/csr.hpp>
 #include <raft/sparse/op/sort.hpp>
 
diff --git a/cpp/src/hdbscan/detail/membership.cuh b/cpp/src/hdbscan/detail/membership.cuh
index 55f4b72cd1..8170a1dc77 100644
--- a/cpp/src/hdbscan/detail/membership.cuh
+++ b/cpp/src/hdbscan/detail/membership.cuh
@@ -28,7 +28,7 @@
 
 #include <cuml/cluster/hdbscan.hpp>
 
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <algorithm>
 
diff --git a/cpp/src/hdbscan/detail/reachability.cuh b/cpp/src/hdbscan/detail/reachability.cuh
index 26080dee8f..eb6c4ba926 100644
--- a/cpp/src/hdbscan/detail/reachability.cuh
+++ b/cpp/src/hdbscan/detail/reachability.cuh
@@ -21,8 +21,6 @@
 #include <raft/cuda_utils.cuh>
 #include <raft/cudart_utils.h>
 
-#include <raft/mr/device/buffer.hpp>
-
 #include <raft/linalg/unary_op.cuh>
 
 #include <raft/sparse/convert/csr.hpp>
diff --git a/cpp/src/hdbscan/detail/select.cuh b/cpp/src/hdbscan/detail/select.cuh
index 4aeb3a5640..aec1c05d8a 100644
--- a/cpp/src/hdbscan/detail/select.cuh
+++ b/cpp/src/hdbscan/detail/select.cuh
@@ -28,7 +28,7 @@
 
 #include <cuml/cluster/hdbscan.hpp>
 
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <algorithm>
 
diff --git a/cpp/src/hdbscan/detail/stabilities.cuh b/cpp/src/hdbscan/detail/stabilities.cuh
index ce443d6867..417b36381e 100644
--- a/cpp/src/hdbscan/detail/stabilities.cuh
+++ b/cpp/src/hdbscan/detail/stabilities.cuh
@@ -28,7 +28,7 @@
 
 #include <cuml/cluster/hdbscan.hpp>
 
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <algorithm>
 
diff --git a/cpp/src/hdbscan/detail/utils.h b/cpp/src/hdbscan/detail/utils.h
index 6db1974fd5..5fa84006d4 100644
--- a/cpp/src/hdbscan/detail/utils.h
+++ b/cpp/src/hdbscan/detail/utils.h
@@ -16,8 +16,6 @@
 
 #pragma once
 
-#include <label/classlabels.cuh>
-
 #include <cub/cub.cuh>
 
 #include <raft/cudart_utils.h>
@@ -27,7 +25,7 @@
 
 #include <cuml/cluster/hdbscan.hpp>
 
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 
 #include <algorithm>
 
diff --git a/cpp/src/hierarchy/pw_dist_graph.cuh b/cpp/src/hierarchy/pw_dist_graph.cuh
index dca5f100ba..573a48f3bf 100644
--- a/cpp/src/hierarchy/pw_dist_graph.cuh
+++ b/cpp/src/hierarchy/pw_dist_graph.cuh
@@ -25,7 +25,6 @@
 #include <rmm/exec_policy.hpp>
 
 #include <raft/linalg/distance_type.h>
-#include <raft/mr/device/buffer.hpp>
 
 // TODO: Not a good strategy for pluggability but will be
 // removed once our dense pairwise distance API is in RAFT
diff --git a/cpp/src/holtwinters/internal/hw_eval.cuh b/cpp/src/holtwinters/internal/hw_eval.cuh
index 8fa118ec0e..0ca25b77f2 100644
--- a/cpp/src/holtwinters/internal/hw_eval.cuh
+++ b/cpp/src/holtwinters/internal/hw_eval.cuh
@@ -20,7 +20,6 @@
 
 #include <raft/cudart_utils.h>
 #include <raft/handle.hpp>
-#include <raft/mr/device/buffer.hpp>
 
 template <typename Dtype>
 __device__ Dtype holtwinters_eval_device(int tid,
diff --git a/cpp/src/holtwinters/internal/hw_optim.cuh b/cpp/src/holtwinters/internal/hw_optim.cuh
index 96f6a1e24d..a105a8730f 100644
--- a/cpp/src/holtwinters/internal/hw_optim.cuh
+++ b/cpp/src/holtwinters/internal/hw_optim.cuh
@@ -20,7 +20,6 @@
 #include "hw_utils.cuh"
 
 #include <raft/cudart_utils.h>
-#include <raft/mr/device/buffer.hpp>
 
 template <typename Dtype>
 __device__ Dtype golden_step(Dtype a, Dtype b, Dtype c)
diff --git a/cpp/src/knn/knn.cu b/cpp/src/knn/knn.cu
index 1c8d2f57dc..381459f7f3 100644
--- a/cpp/src/knn/knn.cu
+++ b/cpp/src/knn/knn.cu
@@ -16,7 +16,7 @@
 
 #include <cuda_runtime.h>
 #include <raft/cuda_utils.cuh>
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 #include <raft/spatial/knn/ann.hpp>
 #include <raft/spatial/knn/ball_cover.hpp>
 
diff --git a/cpp/src/svm/svc_impl.cuh b/cpp/src/svm/svc_impl.cuh
index 9a55a95325..45b8bcd90d 100644
--- a/cpp/src/svm/svc_impl.cuh
+++ b/cpp/src/svm/svc_impl.cuh
@@ -30,7 +30,7 @@
 #include <cuml/svm/svm_parameter.h>
 #include <label/classlabels.cuh>
 #include <matrix/kernelfactory.cuh>
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 #include <raft/linalg/cublas_wrappers.h>
 #include <raft/linalg/unary_op.cuh>
 #include <raft/matrix/matrix.hpp>
diff --git a/cpp/src/tsne/fft_tsne.cuh b/cpp/src/tsne/fft_tsne.cuh
index 05c00b6d50..6b3453da77 100644
--- a/cpp/src/tsne/fft_tsne.cuh
+++ b/cpp/src/tsne/fft_tsne.cuh
@@ -30,7 +30,6 @@
 #include <cufft_utils.h>
 #include <linalg/init.h>
 #include <raft/linalg/eltwise.cuh>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/stats/sum.hpp>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
diff --git a/cpp/src_prims/functions/hinge.cuh b/cpp/src_prims/functions/hinge.cuh
index e7faf31ad4..42b1d400c9 100644
--- a/cpp/src_prims/functions/hinge.cuh
+++ b/cpp/src_prims/functions/hinge.cuh
@@ -28,7 +28,6 @@
 #include <raft/linalg/unary_op.cuh>
 #include <raft/matrix/math.hpp>
 #include <raft/matrix/matrix.hpp>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/stats/mean.hpp>
 #include <raft/stats/sum.hpp>
 #include <rmm/device_uvector.hpp>
diff --git a/cpp/src_prims/functions/linearReg.cuh b/cpp/src_prims/functions/linearReg.cuh
index 02653cb3b3..460250b241 100644
--- a/cpp/src_prims/functions/linearReg.cuh
+++ b/cpp/src_prims/functions/linearReg.cuh
@@ -26,7 +26,6 @@
 #include <raft/linalg/transpose.h>
 #include <raft/matrix/math.hpp>
 #include <raft/matrix/matrix.hpp>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/stats/mean.hpp>
 #include <raft/stats/sum.hpp>
 #include <rmm/device_uvector.hpp>
diff --git a/cpp/src_prims/label/classlabels.cuh b/cpp/src_prims/label/classlabels.cuh
index bd4bd2d7b6..f75975b52e 100644
--- a/cpp/src_prims/label/classlabels.cuh
+++ b/cpp/src_prims/label/classlabels.cuh
@@ -21,7 +21,7 @@
 #include <raft/cuda_utils.cuh>
 #include <raft/cudart_utils.h>
 #include <raft/handle.hpp>
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 #include <raft/linalg/unary_op.cuh>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
diff --git a/cpp/src_prims/linalg/lstsq.cuh b/cpp/src_prims/linalg/lstsq.cuh
index d4810f4ae7..d095f71c10 100644
--- a/cpp/src_prims/linalg/lstsq.cuh
+++ b/cpp/src_prims/linalg/lstsq.cuh
@@ -31,7 +31,6 @@
 #include <raft/linalg/transpose.h>
 #include <raft/matrix/math.hpp>
 #include <raft/matrix/matrix.hpp>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/random/rng.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_scalar.hpp>
diff --git a/cpp/src_prims/linalg/rsvd.cuh b/cpp/src_prims/linalg/rsvd.cuh
index 55a3ecfa86..5a37cbfbb2 100644
--- a/cpp/src_prims/linalg/rsvd.cuh
+++ b/cpp/src_prims/linalg/rsvd.cuh
@@ -26,7 +26,6 @@
 #include <raft/linalg/transpose.h>
 #include <raft/matrix/math.hpp>
 #include <raft/matrix/matrix.hpp>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/random/rng.hpp>
 
 namespace MLCommon {
diff --git a/cpp/src_prims/random/make_regression.cuh b/cpp/src_prims/random/make_regression.cuh
index 21f70170b9..1c5364e6a9 100644
--- a/cpp/src_prims/random/make_regression.cuh
+++ b/cpp/src_prims/random/make_regression.cuh
@@ -31,7 +31,6 @@
 #include <raft/linalg/qr.cuh>
 #include <raft/linalg/transpose.h>
 #include <raft/matrix/matrix.hpp>
-#include <raft/mr/device/buffer.hpp>
 #include <raft/random/rng.hpp>
 #include <rmm/device_uvector.hpp>
 
diff --git a/cpp/test/prims/knn_classify.cu b/cpp/test/prims/knn_classify.cu
index 366ec7b25d..3ec039e1ea 100644
--- a/cpp/test/prims/knn_classify.cu
+++ b/cpp/test/prims/knn_classify.cu
@@ -19,7 +19,7 @@
 #include <iostream>
 #include <raft/cuda_utils.cuh>
 #include <raft/cudart_utils.h>
-#include <raft/label/classlabels.cuh>
+#include <raft/label/classlabels.hpp>
 #include <raft/spatial/knn/knn.hpp>
 #include <random/make_blobs.cuh>
 #include <rmm/device_uvector.hpp>

From 5902b07fa9234068612de4b1a91d03dab27cd0b0 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 1 Feb 2022 18:34:13 -0500
Subject: [PATCH 2/5] Updating raft pin

---
 cpp/cmake/thirdparty/get_raft.cmake | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake
index 63f795d519..519925cf1a 100644
--- a/cpp/cmake/thirdparty/get_raft.cmake
+++ b/cpp/cmake/thirdparty/get_raft.cmake
@@ -57,8 +57,8 @@ set(CUML_BRANCH_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION          ${CUML_MIN_VERSION_raft}
-                        FORK             rapidsai
-                        PINNED_TAG       branch-${CUML_BRANCH_VERSION_raft}
+                        FORK             cjnolet
+                        PINNED_TAG       imp-2204_remaining_public_api_bits
                         USE_RAFT_NN      ${CUML_USE_RAFT_NN}
                         USE_FAISS_STATIC ${CUML_USE_FAISS_STATIC}
                         )

From fc41d8fa02e4a55ae3cc6f7a4f8361609fcbb692 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 1 Feb 2022 18:48:32 -0500
Subject: [PATCH 3/5] Updating copyrights

---
 cpp/src/hdbscan/detail/utils.h            | 2 +-
 cpp/src/holtwinters/internal/hw_eval.cuh  | 2 +-
 cpp/src/holtwinters/internal/hw_optim.cuh | 2 +-
 cpp/src/knn/knn.cu                        | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/src/hdbscan/detail/utils.h b/cpp/src/hdbscan/detail/utils.h
index 5fa84006d4..b5fe03a671 100644
--- a/cpp/src/hdbscan/detail/utils.h
+++ b/cpp/src/hdbscan/detail/utils.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/cpp/src/holtwinters/internal/hw_eval.cuh b/cpp/src/holtwinters/internal/hw_eval.cuh
index 0ca25b77f2..93592811e0 100644
--- a/cpp/src/holtwinters/internal/hw_eval.cuh
+++ b/cpp/src/holtwinters/internal/hw_eval.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/cpp/src/holtwinters/internal/hw_optim.cuh b/cpp/src/holtwinters/internal/hw_optim.cuh
index a105a8730f..0036f76fe2 100644
--- a/cpp/src/holtwinters/internal/hw_optim.cuh
+++ b/cpp/src/holtwinters/internal/hw_optim.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
diff --git a/cpp/src/knn/knn.cu b/cpp/src/knn/knn.cu
index 381459f7f3..e9586c53dd 100644
--- a/cpp/src/knn/knn.cu
+++ b/cpp/src/knn/knn.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

From ec44e7cfa6f2ef76d4f30c3a88d62ab8caa48079 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Feb 2022 19:46:44 -0500
Subject: [PATCH 4/5] Fixing raft::spasrse calls

---
 cpp/src/glm/qn/simple_mat/sparse.hpp | 64 ++++++++++++++--------------
 1 file changed, 33 insertions(+), 31 deletions(-)

diff --git a/cpp/src/glm/qn/simple_mat/sparse.hpp b/cpp/src/glm/qn/simple_mat/sparse.hpp
index 29ad5c8f5c..6695732dc8 100644
--- a/cpp/src/glm/qn/simple_mat/sparse.hpp
+++ b/cpp/src/glm/qn/simple_mat/sparse.hpp
@@ -30,6 +30,8 @@
 #include <raft/mr/device/allocator.hpp>
 #include <rmm/device_uvector.hpp>
 
+#include <raft/sparse/detail/cusparse_wrappers.h>
+
 namespace ML {
 
 /**
@@ -94,7 +96,7 @@ struct SimpleSparseMat : SimpleMat<T> {
     // to swap arguments A and B in cusparseSpMM.
     cusparseDnMatDescr_t descrC;
     auto order = C.ord == COL_MAJOR ? CUSPARSE_ORDER_ROW : CUSPARSE_ORDER_COL;
-    RAFT_CUSPARSE_TRY(raft::sparse::cusparsecreatednmat(
+    RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsecreatednmat(
       &descrC, C.n, C.m, order == CUSPARSE_ORDER_COL ? C.n : C.m, C.data, order));
 
     /*
@@ -115,49 +117,49 @@ struct SimpleSparseMat : SimpleMat<T> {
         ldX'   - leading dimension - m or n, depending on order and transX
      */
     cusparseDnMatDescr_t descrA;
-    RAFT_CUSPARSE_TRY(raft::sparse::cusparsecreatednmat(&descrA,
-                                                        C.ord == A.ord ? A.n : A.m,
-                                                        C.ord == A.ord ? A.m : A.n,
-                                                        A.ord == COL_MAJOR ? A.m : A.n,
-                                                        A.data,
-                                                        order));
+    RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsecreatednmat(&descrA,
+                                                                C.ord == A.ord ? A.n : A.m,
+                                                                C.ord == A.ord ? A.m : A.n,
+                                                                A.ord == COL_MAJOR ? A.m : A.n,
+                                                                A.data,
+                                                                order));
     auto opA =
       transA ^ (C.ord == A.ord) ? CUSPARSE_OPERATION_NON_TRANSPOSE : CUSPARSE_OPERATION_TRANSPOSE;
 
     cusparseSpMatDescr_t descrB;
-    RAFT_CUSPARSE_TRY(
-      raft::sparse::cusparsecreatecsr(&descrB, B.m, B.n, B.nnz, B.row_ids, B.cols, B.values));
+    RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsecreatecsr(
+      &descrB, B.m, B.n, B.nnz, B.row_ids, B.cols, B.values));
     auto opB = transB ? CUSPARSE_OPERATION_NON_TRANSPOSE : CUSPARSE_OPERATION_TRANSPOSE;
 
     auto alg = order == CUSPARSE_ORDER_COL ? CUSPARSE_SPMM_CSR_ALG1 : CUSPARSE_SPMM_CSR_ALG2;
 
     size_t bufferSize;
-    RAFT_CUSPARSE_TRY(raft::sparse::cusparsespmm_bufferSize(handle.get_cusparse_handle(),
-                                                            opB,
-                                                            opA,
-                                                            &alpha,
-                                                            descrB,
-                                                            descrA,
-                                                            &beta,
-                                                            descrC,
-                                                            alg,
-                                                            &bufferSize,
-                                                            stream));
+    RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsespmm_bufferSize(handle.get_cusparse_handle(),
+                                                                    opB,
+                                                                    opA,
+                                                                    &alpha,
+                                                                    descrB,
+                                                                    descrA,
+                                                                    &beta,
+                                                                    descrC,
+                                                                    alg,
+                                                                    &bufferSize,
+                                                                    stream));
 
     raft::interruptible::synchronize(stream);
     rmm::device_uvector<T> tmp(bufferSize, stream);
 
-    RAFT_CUSPARSE_TRY(raft::sparse::cusparsespmm(handle.get_cusparse_handle(),
-                                                 opB,
-                                                 opA,
-                                                 &alpha,
-                                                 descrB,
-                                                 descrA,
-                                                 &beta,
-                                                 descrC,
-                                                 alg,
-                                                 tmp.data(),
-                                                 stream));
+    RAFT_CUSPARSE_TRY(raft::sparse::detail::cusparsespmm(handle.get_cusparse_handle(),
+                                                         opB,
+                                                         opA,
+                                                         &alpha,
+                                                         descrB,
+                                                         descrA,
+                                                         &beta,
+                                                         descrC,
+                                                         alg,
+                                                         tmp.data(),
+                                                         stream));
 
     RAFT_CUSPARSE_TRY(cusparseDestroyDnMat(descrA));
     RAFT_CUSPARSE_TRY(cusparseDestroySpMat(descrB));

From bb3fad45e629027b47f224ba8f57f77281e8c406 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Feb 2022 22:20:17 -0500
Subject: [PATCH 5/5] Reverting raft pin

---
 cpp/cmake/thirdparty/get_raft.cmake | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake
index 519925cf1a..63f795d519 100644
--- a/cpp/cmake/thirdparty/get_raft.cmake
+++ b/cpp/cmake/thirdparty/get_raft.cmake
@@ -57,8 +57,8 @@ set(CUML_BRANCH_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION          ${CUML_MIN_VERSION_raft}
-                        FORK             cjnolet
-                        PINNED_TAG       imp-2204_remaining_public_api_bits
+                        FORK             rapidsai
+                        PINNED_TAG       branch-${CUML_BRANCH_VERSION_raft}
                         USE_RAFT_NN      ${CUML_USE_RAFT_NN}
                         USE_FAISS_STATIC ${CUML_USE_FAISS_STATIC}
                         )