[BUG] CAGRA computation results will go wrong about raft HOT 3 OPEN

import torch
import numpy as np
import cupy as cp

from pylibraft.common import DeviceResources
from pylibraft.neighbors import cagra

import rmm
pool = rmm.mr.PoolMemoryResource(
    rmm.mr.CudaMemoryResource(),
    initial_pool_size=2**30,
    maximum_pool_size=2**32
)
rmm.mr.set_current_device_resource(pool)
rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource())

batch_size = 1000
n_samples = 10
n_features = 64
# n_queries = 1000
k = 3
dataset = torch.randn(batch_size * n_samples, n_features, dtype=torch.float32).cuda()

handle = DeviceResources()
build_params = cagra.IndexParams(metric="sqeuclidean", build_algo='nn_descent', intermediate_graph_degree=32, graph_degree=32,)

seperate_constant = torch.arange(batch_size).view(-1, 1).repeat(1, n_samples).view(-1, 1).cuda() * 1000
# seperate_constant looks like [0, 0, 0,..., 1000, 1000, 1000,..., 2000, 2000, 2000, ...]
dataset += seperate_constant

index = cagra.build(build_params, dataset, handle=handle)
distances, neighbors = cagra.search(cagra.SearchParams(),
                                     index, dataset,
                                     k, handle=handle)
handle.sync()

neighbors = torch.as_tensor(cp.asarray(neighbors).astype(cp.int64), device='cuda')


for i in range(batch_size):
    print(i, sum( (neighbors[i*n_samples:(i+1)*n_samples]<i*n_samples) | (neighbors[i*n_samples:(i+1)*n_samples]>=(i+1)*n_samples)  ).sum())

import torch
import numpy as np
import cupy as cp

from pylibraft.common import DeviceResources
from pylibraft.neighbors import cagra

import rmm
pool = rmm.mr.PoolMemoryResource(
    rmm.mr.CudaMemoryResource(),
    initial_pool_size=2**30,
    maximum_pool_size=2**32
)
rmm.mr.set_current_device_resource(pool)
rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource())

batch_size = 10
n_samples = 1000
n_features = 64
# n_queries = 1000
k = 3
key = torch.randn(batch_size * n_samples, n_features, dtype=torch.float32).cuda()
query = torch.randn(batch_size * n_samples, n_features, dtype=torch.float32).cuda()

handle = DeviceResources()
build_params = cagra.IndexParams(metric="sqeuclidean", build_algo='nn_descent', intermediate_graph_degree=128, graph_degree=128,)

index = cagra.build(build_params, key, handle=handle)
distances, neighbors = cagra.search(cagra.SearchParams(),
                                     index, query,
                                     k, handle=handle)
handle.sync()

neighbors = torch.as_tensor(cp.asarray(neighbors).astype(cp.int64), device='cuda')
distances = torch.as_tensor(distances, device='cuda')

# check if the neighbors are correct
distances_check = torch.zeros(batch_size * n_samples, k, dtype=torch.float32).cuda()
neighbors_check = torch.zeros(batch_size * n_samples, k, dtype=torch.int64).cuda()

for i in range(batch_size):
    build_params = cagra.IndexParams(metric="sqeuclidean", build_algo='nn_descent', intermediate_graph_degree=32, graph_degree=32,)
    index = cagra.build(build_params, key, handle=handle)
    distances_index, neighbors_index = cagra.search(cagra.SearchParams(),
                                         index, query[i*n_samples:(i+1)*n_samples],
                                         k, handle=handle)
    handle.sync()
    distances_check[i*n_samples:(i+1)*n_samples] = torch.as_tensor(distances_index, device='cuda')
    neighbors_check[i*n_samples:(i+1)*n_samples] = torch.as_tensor(cp.asarray(neighbors_index).astype(cp.int64), device='cuda')

    print(i, (distances[i*n_samples:(i+1)*n_samples] - distances_check[i*n_samples:(i+1)*n_samples]).sum())
    print(i, (neighbors[i*n_samples:(i+1)*n_samples] - neighbors_check[i*n_samples:(i+1)*n_samples]).sum())