From 2db9d3d2497dd0a032af7a5e92679cf7c7df696c Mon Sep 17 00:00:00 2001
From: Huy Nguyen <huuanhhuyn@nvidia.com>
Date: Tue, 30 Jun 2026 09:05:04 +0000
Subject: [PATCH 1/2] Fix gather index type in ann_ivf_flat test

---
 cpp/tests/neighbors/ann_ivf_flat.cuh | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/tests/neighbors/ann_ivf_flat.cuh b/cpp/tests/neighbors/ann_ivf_flat.cuh
index ccd7e2a253..c2a9310703 100644
--- a/cpp/tests/neighbors/ann_ivf_flat.cuh
+++ b/cpp/tests/neighbors/ann_ivf_flat.cuh
@@ -311,10 +311,10 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
 
         raft::matrix::gather(
           handle_,
-          raft::make_device_matrix_view<const DataT, uint32_t>(
-            (const DataT*)database.data(), static_cast<uint32_t>(ps.num_db_vecs), idx.dim()),
-          raft::make_device_vector_view<const IdxT, uint32_t>((const IdxT*)list_inds.data_handle(),
-                                                              list_size),
+          raft::make_device_matrix_view<const DataT, int32_t>(
+            (const DataT*)database.data(), static_cast<int32_t>(ps.num_db_vecs), idx.dim()),
+          raft::make_device_vector_view<const IdxT, int32_t>((const IdxT*)list_inds.data_handle(),
+                                                             list_size),
           flat_codes.view());
 
         helpers::codepacker::pack(

From 95e5f429335a600cbc4935b02b572fdd437e5f0d Mon Sep 17 00:00:00 2001
From: Huy Nguyen <huuanhhuyn@nvidia.com>
Date: Tue, 30 Jun 2026 10:11:25 +0000
Subject: [PATCH 2/2] Change flat_codes index type to int

---
 cpp/tests/neighbors/ann_ivf_flat.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/tests/neighbors/ann_ivf_flat.cuh b/cpp/tests/neighbors/ann_ivf_flat.cuh
index c2a9310703..13a3506188 100644
--- a/cpp/tests/neighbors/ann_ivf_flat.cuh
+++ b/cpp/tests/neighbors/ann_ivf_flat.cuh
@@ -307,7 +307,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         auto& list_inds           = extend_index.lists()[label]->indices;
 
         // fetch the flat codes
-        auto flat_codes = raft::make_device_matrix<DataT, uint32_t>(handle_, list_size, idx.dim());
+        auto flat_codes = raft::make_device_matrix<DataT, int32_t>(handle_, list_size, idx.dim());
 
         raft::matrix::gather(
           handle_,