Skip to content

Commit afb2b3f

Browse files
jinsolpjinsolp
authored
Expose supported brute force metrics in all_neighbors (#1827)
Closes #1715 Exposes supported brute force metrics in `all_neighbors` API. Authors: - Jinsol Park (https://github.com/jinsolp) Approvers: - Victor Lafargue (https://github.com/viclafargue) - Tarang Jain (https://github.com/tarang-jain) - Corey J. Nolet (https://github.com/cjnolet) URL: #1827
1 parent 726ba20 commit afb2b3f

3 files changed

Lines changed: 89 additions & 35 deletions

File tree

‎cpp/src/neighbors/all_neighbors/all_neighbors.cuh‎

Lines changed: 32 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,52 +9,60 @@
99
#include <cuvs/neighbors/all_neighbors.hpp>
1010
#include <raft/matrix/shift.cuh>
1111
#include <raft/util/cudart_utils.hpp>
12+
#include <unordered_set>
1213

1314
namespace cuvs::neighbors::all_neighbors::detail {
1415
using namespace cuvs::neighbors;
1516

1617
GRAPH_BUILD_ALGO check_params_validity(const all_neighbors_params& params,
1718
bool do_mutual_reachability_dist)
1819
{
20+
using DT = cuvs::distance::DistanceType;
21+
22+
// InnerProduct is not supported for mutual reachability distance, because mutual reachability
23+
// distance takes "max" of core distances and pairwise distance.
24+
static const std::unordered_set<DT> mrd_allowed_metrics = {
25+
DT::L2Expanded, DT::L2SqrtExpanded, DT::CosineExpanded};
26+
27+
static const std::unordered_set<DT> bf_allowed_metrics = {DT::L2Expanded,
28+
DT::L2SqrtExpanded,
29+
DT::CosineExpanded,
30+
DT::L1,
31+
DT::L2Unexpanded,
32+
DT::L2SqrtUnexpanded,
33+
DT::InnerProduct,
34+
DT::Linf,
35+
DT::Canberra,
36+
DT::LpUnexpanded,
37+
DT::CorrelationExpanded,
38+
DT::JensenShannon};
39+
40+
static const std::unordered_set<DT> nnd_allowed_metrics = {
41+
DT::L2Expanded, DT::L2SqrtExpanded, DT::CosineExpanded, DT::InnerProduct};
42+
1943
if (std::holds_alternative<graph_build_params::brute_force_params>(params.graph_build_params)) {
2044
if (do_mutual_reachability_dist) {
21-
// InnerProduct is not supported for mutual reachability distance, because mutual reachability
22-
// distance takes "max" of core distances and pairwise distance.
23-
auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded ||
24-
params.metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
25-
params.metric == cuvs::distance::DistanceType::CosineExpanded;
2645
RAFT_EXPECTS(
27-
allowed_metrics,
46+
mrd_allowed_metrics.count(params.metric),
2847
"Distance metric for all-neighbors build with brute force for computing mutual "
2948
"reachability distance should be L2Expanded, L2SqrtExpanded, or CosineExpanded.");
3049
} else {
31-
auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded ||
32-
params.metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
33-
params.metric == cuvs::distance::DistanceType::CosineExpanded ||
34-
params.metric == cuvs::distance::DistanceType::InnerProduct;
35-
RAFT_EXPECTS(allowed_metrics,
36-
"Distance metric for all-neighbors build with brute force should be L2Expanded, "
37-
"L2SqrtExpanded, CosineExpanded, or InnerProduct.");
50+
RAFT_EXPECTS(
51+
bf_allowed_metrics.count(params.metric),
52+
"Distance metric for all-neighbors build with brute force should be L2Expanded, "
53+
"L2SqrtExpanded, CosineExpanded, L1, L2Unexpanded, L2SqrtUnexpanded, InnerProduct, Linf, "
54+
"Canberra, LpUnexpanded, CorrelationExpanded, or JensenShannon.");
3855
}
3956
return GRAPH_BUILD_ALGO::BRUTE_FORCE;
4057
} else if (std::holds_alternative<graph_build_params::nn_descent_params>(
4158
params.graph_build_params)) {
4259
if (do_mutual_reachability_dist) {
43-
// InnerProduct is not supported for mutual reachability distance, because mutual reachability
44-
// distance takes "max" of core distances and pairwise distance.
45-
auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded ||
46-
params.metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
47-
params.metric == cuvs::distance::DistanceType::CosineExpanded;
4860
RAFT_EXPECTS(
49-
allowed_metrics,
61+
mrd_allowed_metrics.count(params.metric),
5062
"Distance metric for all-neighbors build with NN Descent for computing mutual reachability "
5163
"distance should be L2Expanded, L2SqrtExpanded, or CosineExpanded.");
5264
} else {
53-
auto allowed_metrics = params.metric == cuvs::distance::DistanceType::L2Expanded ||
54-
params.metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
55-
params.metric == cuvs::distance::DistanceType::CosineExpanded ||
56-
params.metric == cuvs::distance::DistanceType::InnerProduct;
57-
RAFT_EXPECTS(allowed_metrics,
65+
RAFT_EXPECTS(nnd_allowed_metrics.count(params.metric),
5866
"Distance metric for all-neighbors build with NN Descent should be L2Expanded, "
5967
"L2SqrtExpanded, CosineExpanded, or InnerProduct.");
6068
}

‎python/cuvs/cuvs/neighbors/all_neighbors/all_neighbors.pyx‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#
2-
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
2+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
33
# SPDX-License-Identifier: Apache-2.0
44
#
55
# cython: language_level=3
@@ -111,11 +111,12 @@ cdef class AllNeighborsParams:
111111
)
112112

113113
# Check metric consistency
114-
ivf_pq_metric = ivf_pq_params.metric
115-
if ivf_pq_metric != metric:
114+
metric_type = DISTANCE_TYPES[metric]
115+
ivf_pq_metric_type = DISTANCE_TYPES[ivf_pq_params.metric]
116+
if ivf_pq_metric_type != metric_type:
116117
raise ValueError(
117118
f"Metric conflict: AllNeighborsParams metric '{metric}' "
118-
f"does not match IVF-PQ metric '{ivf_pq_metric}'. Please "
119+
f"does not match IVF-PQ metric '{ivf_pq_params.metric}'. Please "
119120
f"ensure both use the same metric."
120121
)
121122

@@ -127,11 +128,12 @@ cdef class AllNeighborsParams:
127128
)
128129

129130
# Check metric consistency
130-
nn_descent_metric = nn_descent_params.metric
131-
if nn_descent_metric != metric:
131+
metric_type = DISTANCE_TYPES[metric]
132+
nn_descent_metric_type = DISTANCE_TYPES[nn_descent_params.metric]
133+
if nn_descent_metric_type != metric_type:
132134
raise ValueError(
133135
f"Metric conflict: AllNeighborsParams metric '{metric}' "
134-
f"does not match NN-Descent metric '{nn_descent_metric}'. "
136+
f"does not match NN-Descent metric '{nn_descent_params.metric}'. "
135137
f"Please ensure both use the same metric."
136138
)
137139

‎python/cuvs/cuvs/tests/test_all_neighbors.py‎

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
22
# SPDX-License-Identifier: Apache-2.0
33
#
44

@@ -36,15 +36,32 @@ def make_cosine(
3636

3737
@pytest.mark.parametrize("algo", ["nn_descent", "brute_force", "ivf_pq"])
3838
@pytest.mark.parametrize("cluster", ["single_cluster", "multi_cluster"])
39-
@pytest.mark.parametrize("metric", ["sqeuclidean", "cosine"])
39+
@pytest.mark.parametrize(
40+
"metric",
41+
[
42+
"sqeuclidean",
43+
"l2",
44+
"cosine",
45+
"l1",
46+
"inner_product",
47+
"chebyshev",
48+
"canberra",
49+
"minkowski",
50+
"correlation",
51+
"jensenshannon",
52+
],
53+
)
4054
def test_all_neighbors_device_build_quality(algo, cluster, metric):
4155
"""Test device build with quality validation against brute force ground
4256
truth.
4357
"""
4458
n_rows, n_cols, k = 7151, 64, 16
4559

46-
if algo == "ivf_pq" and metric == "cosine":
47-
pytest.skip("Skipping IVF-PQ with cosine distance")
60+
ivf_pq_valid_metrics = {"sqeuclidean"}
61+
nnd_valid_metrics = {"sqeuclidean", "l2", "cosine", "inner_product"}
62+
is_invalid = (algo == "ivf_pq" and metric not in ivf_pq_valid_metrics) or (
63+
algo == "nn_descent" and metric not in nnd_valid_metrics
64+
)
4865

4966
if cluster == "single_cluster":
5067
overlap_factor = 0
@@ -57,6 +74,21 @@ def test_all_neighbors_device_build_quality(algo, cluster, metric):
5774
X, _ = make_cosine(
5875
n_samples=n_rows, n_features=n_cols, random_state=42
5976
)
77+
elif metric == "jensenshannon":
78+
# Jensen-Shannon requires non-negative values representing probability distributions
79+
X, _ = make_blobs(
80+
n_samples=n_rows,
81+
n_features=n_cols,
82+
centers=10,
83+
cluster_std=1.0,
84+
center_box=(0.0, 10.0), # Non-negative values only
85+
random_state=42,
86+
)
87+
# Normalize each row to sum to 1 (probability distribution)
88+
X = np.abs(X) # Ensure non-negative
89+
row_sums = X.sum(axis=1, keepdims=True)
90+
row_sums[row_sums == 0] = 1 # Avoid division by zero
91+
X = X / row_sums
6092
else:
6193
X, _ = make_blobs(
6294
n_samples=n_rows,
@@ -98,6 +130,18 @@ def test_all_neighbors_device_build_quality(algo, cluster, metric):
98130
)
99131

100132
res = Resources()
133+
134+
if is_invalid:
135+
with pytest.raises(Exception, match="Distance metric"):
136+
all_neighbors.build(
137+
X_device,
138+
k,
139+
params,
140+
distances=cupy.empty((n_rows, k), dtype=cupy.float32),
141+
resources=res,
142+
)
143+
return
144+
101145
indices, distances = all_neighbors.build(
102146
X_device,
103147
k,

0 commit comments

Comments
 (0)