def search_index_pytorch(index, x, k, D=None, I=None):
"""call the search function of an index with pytorch tensor I/O (CPU
and GPU supported)"""
assert x.is_contiguous()
n, d = x.size()
assert d == index.d
if D is None:
if x.is_cuda:
D = torch.cuda.FloatTensor(n, k)
else:
D = torch.FloatTensor(n, k)
else:
assert D.__class__ in (torch.FloatTensor, torch.cuda.FloatTensor)
assert D.size() == (n, k)
assert D.is_contiguous()
if I is None:
if x.is_cuda:
I = torch.cuda.LongTensor(n, k)
else:
I = torch.LongTensor(n, k)
else:
assert I.__class__ in (torch.LongTensor, torch.cuda.LongTensor)
assert I.size() == (n, k)
assert I.is_contiguous()
xptr = x.storage().data_ptr()
Iptr = I.storage().data_ptr()
Dptr = D.storage().data_ptr()
index.search_c(n, faiss.cast_integer_to_float_ptr(xptr),
k, faiss.cast_integer_to_float_ptr(Dptr),
faiss.cast_integer_to_long_ptr(Iptr))
return D, I
def swig_ptr_from_FloatTensor(x):
assert x.is_contiguous()
assert x.dtype == torch.float32
return faiss.cast_integer_to_float_ptr(x.storage().data_ptr() +
x.storage_offset() * 4)
def _swig_ptr_from_cupy_float32_array(x): # pragma: no cover
assert x.flags.c_contiguous
assert x.dtype == np.float32
data_ptr = x.__cuda_array_interface__['data'][0]
return faiss.cast_integer_to_float_ptr(data_ptr)
def swig_ptr_from_FloatTensor(x):
""" gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """
assert x.is_contiguous()
assert x.dtype == torch.float32
return faiss.cast_integer_to_float_ptr(
x.storage().data_ptr() + x.storage_offset() * 4)
