diff --git a/faiss/utils/distances.cpp b/faiss/utils/distances.cpp index b50231b2b0..a94301a42f 100644 --- a/faiss/utils/distances.cpp +++ b/faiss/utils/distances.cpp @@ -786,9 +786,11 @@ void fvec_inner_products_by_idx( const float* xj = x + j * d; float* __restrict ipj = ip + j * ny; for (size_t i = 0; i < ny; i++) { - if (idsj[i] < 0) - continue; - ipj[i] = fvec_inner_product(xj, y + d * idsj[i], d); + if (idsj[i] < 0) { + ipj[i] = -INFINITY; + } else { + ipj[i] = fvec_inner_product(xj, y + d * idsj[i], d); + } } } } @@ -809,9 +811,11 @@ void fvec_L2sqr_by_idx( const float* xj = x + j * d; float* __restrict disj = dis + j * ny; for (size_t i = 0; i < ny; i++) { - if (idsj[i] < 0) - continue; - disj[i] = fvec_L2sqr(xj, y + d * idsj[i], d); + if (idsj[i] < 0) { + disj[i] = INFINITY; + } else { + disj[i] = fvec_L2sqr(xj, y + d * idsj[i], d); + } } } } @@ -828,6 +832,8 @@ void pairwise_indexed_L2sqr( for (int64_t j = 0; j < n; j++) { if (ix[j] >= 0 && iy[j] >= 0) { dis[j] = fvec_L2sqr(x + d * ix[j], y + d * iy[j], d); + } else { + dis[j] = INFINITY; } } } @@ -844,6 +850,8 @@ void pairwise_indexed_inner_product( for (int64_t j = 0; j < n; j++) { if (ix[j] >= 0 && iy[j] >= 0) { dis[j] = fvec_inner_product(x + d * ix[j], y + d * iy[j], d); + } else { + dis[j] = -INFINITY; } } } diff --git a/faiss/utils/distances.h b/faiss/utils/distances.h index 06f4edda32..898edeeb0e 100644 --- a/faiss/utils/distances.h +++ b/faiss/utils/distances.h @@ -198,8 +198,16 @@ void fvec_sub(size_t d, const float* a, const float* b, float* c); * Compute a subset of distances ***************************************************************************/ -/* compute the inner product between x and a subset y of ny vectors, - whose indices are given by idy. */ +/** compute the inner product between x and a subset y of ny vectors defined by + * ids + * + * ip(i, j) = inner_product(x(i, :), y(ids(i, j), :)) + * + * @param ip output array, size nx * ny + * @param x first-term vector, size nx * d + * @param y second-term vector, size (max(ids) + 1) * d + * @param ids ids to sample from y, size nx * ny + */ void fvec_inner_products_by_idx( float* ip, const float* x, @@ -209,7 +217,16 @@ void fvec_inner_products_by_idx( size_t nx, size_t ny); -/* same but for a subset in y indexed by idsy (ny vectors in total) */ +/** compute the squared L2 distances between x and a subset y of ny vectors + * defined by ids + * + * dis(i, j) = inner_product(x(i, :), y(ids(i, j), :)) + * + * @param dis output array, size nx * ny + * @param x first-term vector, size nx * d + * @param y second-term vector, size (max(ids) + 1) * d + * @param ids ids to sample from y, size nx * ny + */ void fvec_L2sqr_by_idx( float* dis, const float* x, @@ -236,7 +253,14 @@ void pairwise_indexed_L2sqr( const int64_t* iy, float* dis); -/* same for inner product */ +/** compute dis[j] = inner_product(x[ix[j]], y[iy[j]]) forall j=0..n-1 + * + * @param x size (max(ix) + 1, d) + * @param y size (max(iy) + 1, d) + * @param ix size n + * @param iy size n + * @param dis size n + */ void pairwise_indexed_inner_product( size_t d, size_t n,