def knn(queries, base, nnn=1, distance_type=2, nt=1):
_check_row_float32(base)
_check_row_float32(queries)
n, d = base.shape
nq, d2 = queries.shape
assert d == d2, "base and queries must have same nb of rows (got %d != %d) " % (
d, d2)
idx = numpy.empty((nq, nnn), dtype=numpy.int32)
dis = numpy.empty((nq, nnn), dtype=numpy.float32)
yael.knn_full_thread(distance_type, nq, n, d, nnn, fvec_ref(base),
fvec_ref(queries), None, ivec_ref(idx), fvec_ref(dis),
nt)
return idx, dis
def knn(queries, base,
nnn = 1,
distance_type = 2,
nt = 1):
_check_row_float32(base)
_check_row_float32(queries)
n, d = base.shape
nq, d2 = queries.shape
assert d == d2, "base and queries must have same nb of rows (got %d != %d) " % (d, d2)
idx = numpy.empty((nq, nnn), dtype = numpy.int32)
dis = numpy.empty((nq, nnn), dtype = numpy.float32)
yael.knn_full_thread(distance_type,
nq, n, d, nnn,
yael.numpy_to_fvec_ref(base),
yael.numpy_to_fvec_ref(queries),
None,
yael.numpy_to_ivec_ref(idx),
yael.numpy_to_fvec_ref(dis),
nt)
return idx, dis
