IVF-Flat in FP16 accumulate distances in FP16, which can lead to infinite distances.
To ensure accurate distance computation, FP32 should be used for accumulation.
This can be reproduced by changing the uniform distribution we're using in test to a uniform between 0.0 and 100.0
Python example
from cuvs.neighbors import ivf_flat
import cupy as cp
a=cp.random.uniform(0,100,(1000,128)).astype(cp.float16)
b=cp.random.uniform(0,100,(20,128)).astype(cp.float16)
index = ivf_flat.build(ivf_flat.IndexParams(n_lists=20), a)
distances, neighbors = index.search(ivf_flat.SearchParams(n_probes=20), index, b, 1)
print(distances.copy_to_host())
# [inf, inf, inf, ...]
IVF-Flat in FP16 accumulate distances in FP16, which can lead to infinite distances.
To ensure accurate distance computation, FP32 should be used for accumulation.
This can be reproduced by changing the uniform distribution we're using in test to a uniform between 0.0 and 100.0
Python example