diff --git a/python/cuml/cuml/neighbors/nearest_neighbors.pyx b/python/cuml/cuml/neighbors/nearest_neighbors.pyx index ddf4021f15..c2a2ade762 100644 --- a/python/cuml/cuml/neighbors/nearest_neighbors.pyx +++ b/python/cuml/cuml/neighbors/nearest_neighbors.pyx @@ -192,6 +192,7 @@ void swap_kernel(long long int* I, float* D, int n_rows, int n_cols) { def _drop_self_edges(distances, indices): """Drop edges between a point and itself in the knn graph""" + indices = cp.ascontiguousarray(indices, dtype=cp.int64) rows, cols = indices.shape # Launch config diff --git a/python/cuml/cuml/neighbors/weights.py b/python/cuml/cuml/neighbors/weights.py index ca369c2ed3..1a1f058abd 100644 --- a/python/cuml/cuml/neighbors/weights.py +++ b/python/cuml/cuml/neighbors/weights.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import cupy as cp @@ -56,7 +56,9 @@ def compute_weights(distances, weights): return raw_weights elif callable(weights): # Custom callable weights (raw, not normalized) - raw_weights = cp.asarray(weights(distances), dtype=cp.float32) + raw_weights = cp.asarray( + weights(distances), dtype="float32", order="C" + ) # Return raw weights - normalization will be done in C++ kernel return raw_weights else: diff --git a/python/cuml/tests/test_kneighbors_regressor.py b/python/cuml/tests/test_kneighbors_regressor.py index a0402c8b3b..59b91ef000 100644 --- a/python/cuml/tests/test_kneighbors_regressor.py +++ b/python/cuml/tests/test_kneighbors_regressor.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -173,6 +173,40 @@ def test_weights_predict(weights, n_neighbors): np.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) +def test_callable_weights_non_c_contiguous_cupy_view(): + X, y = make_regression( + n_samples=64, + n_features=6, + n_informative=4, + random_state=42, + ) + X = X.astype(np.float32) + y = y.astype(np.float32) + + X_train, X_test = X[:48], X[48:] + y_train = y[:48] + + def non_c_contiguous_weights(distances): + xp = cp if isinstance(distances, cp.ndarray) else np + weights = xp.asarray(1.0 / (1.0 + distances), dtype=xp.float32) + base = xp.empty((weights.shape[1], weights.shape[0]), dtype=xp.float32) + base[...] = weights.T + result = base.T + assert result.dtype == xp.float32 + assert not result.flags.c_contiguous + return result + + knn_cu = cuKNN(n_neighbors=5, weights=non_c_contiguous_weights) + knn_cu.fit(X_train, y_train) + pred_cu = knn_cu.predict(cp.asarray(X_test)) + + knn_sk = skKNN(n_neighbors=5, weights=non_c_contiguous_weights) + knn_sk.fit(X_train, y_train) + pred_sk = knn_sk.predict(X_test) + + cp.testing.assert_allclose(pred_cu, pred_sk, rtol=1e-4, atol=1e-4) + + @pytest.mark.parametrize("weights", ["uniform", "distance"]) def test_weights_multioutput(weights): """Test weights parameter with multioutput regression.""" diff --git a/python/cuml/tests/test_nearest_neighbors.py b/python/cuml/tests/test_nearest_neighbors.py index 5f72a7dfa0..1cba505752 100644 --- a/python/cuml/tests/test_nearest_neighbors.py +++ b/python/cuml/tests/test_nearest_neighbors.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -712,6 +712,58 @@ def test_nearest_neighbors_sparse( assert (len(diffs[diffs > 0]) / len(np.ravel(skI))) <= 0.005 +def test_nearest_neighbors_sparse_x_none_self_edge_swap_indices_int64(): + X_dense = np.array( + [ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [10.0, 0.0, 0.0, 0.0], + [10.25, 0.0, 0.0, 0.0], + ], + dtype=np.float32, + ) + X = cupyx.scipy.sparse.csr_matrix(cp.asarray(X_dense)) + + nn = cuKNN( + metric="euclidean", + n_neighbors=2, + algorithm="brute", + output_type="cupy", + ) + nn.fit(X) + + explicit_distances, explicit_indices = nn.kneighbors(X, n_neighbors=2) + assert explicit_indices.dtype == cp.int32 + assert explicit_distances.dtype == cp.float32 + + distances, indices = nn.kneighbors(X=None, n_neighbors=1) + assert indices.dtype == cp.int64 + assert indices.flags.c_contiguous + assert distances.dtype == cp.float32 + assert distances.flags.c_contiguous + + indices_np = cp.asnumpy(indices) + distances_np = cp.asnumpy(distances) + for row, row_indices in enumerate(indices_np): + assert row not in row_indices + + sk_distances, sk_indices = ( + skKNN( + metric="euclidean", + n_neighbors=1, + algorithm="brute", + ) + .fit(X_dense) + .kneighbors(X=None, n_neighbors=1) + ) + + np.testing.assert_allclose( + distances_np, sk_distances, atol=1e-5, rtol=1e-5 + ) + for row_indices, expected_indices in zip(indices_np, sk_indices): + assert set(row_indices.tolist()) == set(expected_indices.tolist()) + + @pytest.mark.parametrize("n_neighbors", [1, 5, 6]) def test_haversine(n_neighbors): hoboken_nj = [40.745255, -74.034775]