From c975a9006814585834207922d5bea3ae65567b41 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Tue, 16 Jun 2026 11:04:19 -0500 Subject: [PATCH 1/3] Cleanup `extract_knn_graph` - Moves `extract_knn_graph` to `cuml.manifold.utils`. This function is only used by `cuml.manifold` and is too specific to belong in a general sparse utils module. - Modernizes the code. Applies new input validation, and removes use of legacy `CumlArray`/`SparseCumlArray`. - Fixes bug in handling any KNN graph input, where the order of `.data` was assumed to be correct (only true if a direct output of `kneighbors_graph`), leading to incorrect results. - Fixes bug in handling of KNN graph and tuple input, where self references were assumed to be present, but not validated or documented as required. For the tuple input we now do a cheap sanity check that self references are included, erroring if they're not present. For KNN graph inputs we properly handle all cases (never included, included but sometimes non-zero-distances, and included), and massage the input to satisfy the requirement. We need to do this for KNN inputs since a canonical KNN graph won't include true self references (which should be 0 distances sans rounding errors), but `kneighbors_graph` will include them if `include_self=True` is provided. - Fixes bug in zero-copy handling of KNN graphs. Previously we always copied to device, and also always made additional unnecessary copies. - Improves validation and error messages - Updates docstrings for `precomputed_knn` and `knn_graph` to better indicate the supported options and requirements. - Adds thorough tests for all input types and configurations --- python/cuml/cuml/common/sparsefuncs.py | 215 +----------------- python/cuml/cuml/manifold/t_sne.pyx | 67 +++--- python/cuml/cuml/manifold/umap/umap.pyx | 94 ++++---- python/cuml/cuml/manifold/utils.py | 279 ++++++++++++++++++++++++ python/cuml/tests/test_umap.py | 225 +++++++++++++++++-- 5 files changed, 567 insertions(+), 313 deletions(-) create mode 100644 python/cuml/cuml/manifold/utils.py diff --git a/python/cuml/cuml/common/sparsefuncs.py b/python/cuml/cuml/common/sparsefuncs.py index 3a858ffb25..578ac229d7 100644 --- a/python/cuml/cuml/common/sparsefuncs.py +++ b/python/cuml/cuml/common/sparsefuncs.py @@ -1,20 +1,14 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # - import math import cupy as cp import cupyx import numpy as np -from cupyx.scipy.sparse import coo_matrix as cp_coo_matrix -from cupyx.scipy.sparse import csc_matrix as cp_csc_matrix -from cupyx.scipy.sparse import csr_matrix as cp_csr_matrix -from scipy.sparse import coo_matrix, csc_matrix, csr_matrix from cuml.common.kernel_utils import cuda_kernel_factory -from cuml.internals.input_utils import input_to_cuml_array, input_to_cupy_array def _map_l1_norm_kernel(dtype): @@ -168,210 +162,3 @@ def _insert_zeros(ary, zero_indices): new_ary[data_mask] = ary return new_ary - - -def extract_sparse_knn_graph(knn_graph): - """ - Converts KNN graph from CSR, COO and CSC formats into separate - distance and indice arrays. Input can be a cupy sparse graph (device) - or a numpy sparse graph (host). - - Returns - ------- - tuple or None - (knn_indices, knn_dists, n_samples) where indices and dists are flattened - arrays and n_samples is the number of rows in the graph, or None if - the format is not supported. - """ - if isinstance(knn_graph, (csc_matrix, cp_csc_matrix)): - knn_graph = cupyx.scipy.sparse.csr_matrix(knn_graph) - n_samples = knn_graph.shape[0] - reordering = knn_graph.data.reshape((n_samples, -1)) - reordering = reordering.argsort() - n_neighbors = reordering.shape[1] - reordering += (cp.arange(n_samples) * n_neighbors)[:, np.newaxis] - reordering = reordering.flatten() - knn_graph.indices = knn_graph.indices[reordering] - knn_graph.data = knn_graph.data[reordering] - - knn_indices = None - if isinstance(knn_graph, (csr_matrix, cp_csr_matrix)): - knn_indices = knn_graph.indices - n_samples = knn_graph.shape[0] - elif isinstance(knn_graph, (coo_matrix, cp_coo_matrix)): - knn_indices = knn_graph.col - n_samples = knn_graph.shape[0] - - if knn_indices is not None: - knn_dists = knn_graph.data - return knn_indices, knn_dists, n_samples - else: - return None - - -def extract_pairwise_dists(pw_dists, n_neighbors): - """ - Extract the nearest neighbors distances and indices - from a pairwise distance matrix. - - Parameters - ---------- - pw_dists: paiwise distances matrix of shape (n_samples, n_samples) - n_neighbors: number of nearest neighbors - - (inspired from Scikit-Learn code) - """ - pw_dists, _, _, _ = input_to_cupy_array(pw_dists) - - n_rows = pw_dists.shape[0] - sample_range = cp.arange(n_rows)[:, None] - knn_indices = cp.argpartition(pw_dists, n_neighbors - 1, axis=1) - knn_indices = knn_indices[:, :n_neighbors] - argdist = cp.argsort(pw_dists[sample_range, knn_indices]) - knn_indices = knn_indices[sample_range, argdist] - knn_dists = pw_dists[sample_range, knn_indices] - return knn_indices, knn_dists - - -def _determine_k_from_arrays( - knn_indices_arr, n_neighbors, n_samples_hint=None -): - """Determine k (neighbors per sample) from array shape.""" - if len(knn_indices_arr.shape) == 2: - return knn_indices_arr.shape[1] - - # 1D flattened array - infer n_samples and k - total_elements = knn_indices_arr.shape[0] - n_samples = ( - n_samples_hint - if n_samples_hint is not None - else total_elements // n_neighbors - ) - - if total_elements % n_samples != 0: - raise ValueError( - f"Precomputed KNN data has {total_elements} total elements which is not evenly " - f"divisible by {n_samples} samples. Expected {n_samples * n_neighbors} elements " - f"for n_neighbors={n_neighbors}." - ) - - return total_elements // n_samples - - -def extract_knn_graph(knn_info, n_neighbors, mem_type="device"): - """ - Extract the nearest neighbors distances and indices - from the knn_info parameter. - - Parameters - ---------- - knn_info : array / sparse array / tuple, optional (device or host) - Either one of : - - Tuple (indices, distances) of arrays of - shape (n_samples, n_neighbors) - - Pairwise distances dense array of shape (n_samples, n_samples) - - KNN graph sparse array (preferably CSR/COO) - n_neighbors: number of nearest neighbors - """ - if knn_info is None: - return None - - # Extract indices, distances, and optional n_samples hint - deepcopy = False - n_samples_hint = None - - if isinstance(knn_info, tuple): - knn_indices, knn_dists = knn_info - elif isinstance( - knn_info, - ( - csr_matrix, - coo_matrix, - csc_matrix, - cp_csr_matrix, - cp_coo_matrix, - cp_csc_matrix, - ), - ): - # Sparse matrix - result = extract_sparse_knn_graph(knn_info) - if result is None: - return None - knn_indices, knn_dists, n_samples_hint = result - deepcopy = True - else: - # Dense pairwise distance matrix - result = extract_pairwise_dists(knn_info, n_neighbors) - if result is None: - return None - knn_indices, knn_dists = result - - # Validate the extracted data - knn_indices_arr = ( - knn_indices - if hasattr(knn_indices, "shape") - else np.asarray(knn_indices) - ) - knn_dists_arr = ( - knn_dists if hasattr(knn_dists, "shape") else np.asarray(knn_dists) - ) - - if knn_indices_arr.shape != knn_dists_arr.shape: - raise ValueError( - f"Precomputed KNN indices and distances must have the same shape. " - f"Got indices shape {knn_indices_arr.shape} and distances shape {knn_dists_arr.shape}." - ) - - if len(knn_indices_arr.shape) not in (1, 2): - raise ValueError( - f"Precomputed KNN indices must be 1D or 2D array, got shape {knn_indices_arr.shape}" - ) - - # Determine actual k and validate against expected n_neighbors - k_provided = _determine_k_from_arrays( - knn_indices_arr, n_neighbors, n_samples_hint - ) - - if k_provided < n_neighbors: - raise ValueError( - f"Precomputed KNN data has {k_provided} neighbors per sample, " - f"but n_neighbors={n_neighbors} was specified. " - f"Cannot use fewer neighbors than requested. " - f"Please provide KNN data with at least {n_neighbors} neighbors per sample." - ) - elif k_provided > n_neighbors: - # Trim excess neighbors - if len(knn_indices_arr.shape) == 2: - # 2D array case: trim columns - knn_indices = knn_indices[:, :n_neighbors] - knn_dists = knn_dists[:, :n_neighbors] - else: - # 1D flattened array case: reshape, trim, and flatten - n_samples = knn_indices_arr.shape[0] // k_provided - knn_indices = knn_indices.reshape((n_samples, k_provided))[ - :, :n_neighbors - ] - knn_dists = knn_dists.reshape((n_samples, k_provided))[ - :, :n_neighbors - ] - - # Convert to CumlArray - knn_indices_m, _, _, _ = input_to_cuml_array( - knn_indices.flatten(), - order="C", - deepcopy=deepcopy, - check_dtype=np.int64, - convert_to_dtype=np.int64, - convert_to_mem_type=mem_type, - ) - - knn_dists_m, _, _, _ = input_to_cuml_array( - knn_dists.flatten(), - order="C", - deepcopy=deepcopy, - check_dtype=np.float32, - convert_to_dtype=np.float32, - convert_to_mem_type=mem_type, - ) - - return knn_indices_m, knn_dists_m diff --git a/python/cuml/cuml/manifold/t_sne.pyx b/python/cuml/cuml/manifold/t_sne.pyx index d94661198c..e50af9bf85 100644 --- a/python/cuml/cuml/manifold/t_sne.pyx +++ b/python/cuml/cuml/manifold/t_sne.pyx @@ -8,7 +8,6 @@ import numpy as np from cuml.common.array_descriptor import CumlArrayDescriptor from cuml.common.doc_utils import generate_docstring from cuml.common.sparse_utils import is_sparse -from cuml.common.sparsefuncs import extract_knn_graph from cuml.internals.array import CumlArray from cuml.internals.base import Base, get_handle from cuml.internals.interop import ( @@ -20,6 +19,7 @@ from cuml.internals.interop import ( from cuml.internals.mixins import CMajorInputTagMixin, SparseInputTagMixin from cuml.internals.outputs import reflect from cuml.internals.validation import check_inputs, check_random_seed +from cuml.manifold.utils import extract_knn_graph from libc.stdint cimport int64_t, uintptr_t from libcpp cimport bool @@ -329,14 +329,26 @@ class TSNE(InteropMixin, 'sqeuclidean' metric, the distances will still be squared when True. Note: This argument should likely be set to False for distance metrics other than 'euclidean' and 'l2'. - precomputed_knn : array / sparse array / tuple, optional (device or host) - Either one of a tuple (indices, distances) of - arrays of shape (n_samples, n_neighbors), a pairwise distances - dense array of shape (n_samples, n_samples) or a KNN graph - sparse array (preferably CSR/COO). This feature allows - the precomputation of the KNN outside of TSNE - and also allows the use of a custom distance function. This function - should match the metric used to train the TSNE embeedings. + precomputed_knn : tuple[array, array], sparse-matrix, array, optional + This feature allows the precomputation of the KNN outside of TSNE. + Options are: + + - A tuple (indices, distances) of dense arrays of shape (n_samples, + n_neighbors), where n_neighbors is >= the ``n_neighbors`` parameter. + Self references should be included (i.e. the first column of + `indices` should be [0, 1, ...], denotating that the nearest neighbor + to each row is itself). This is the most efficient representation. + + - A sparse matrix KNN graph, as may be output by + ``cuml.neighbors.kneighbors_graph`` with ``mode="distance"`` and + ``include_self=True``. The ``n_neighbors`` used to calculate the + graph must be >= the ``n_neighbors`` parameter. + + - A pairwise distances dense array of shape (n_samples, n_samples). + + In all cases the KNN should be computed using the same ``metric`` as + provided to ``TSNE``. + output_type : {'input', 'array', 'dataframe', 'series', 'df_obj', \ 'numba', 'cupy', 'numpy', 'cudf', 'pandas'}, default=None Return results and set estimator attributes to the indicated output @@ -562,15 +574,13 @@ class TSNE(InteropMixin, Parameters ---------- - knn_graph : array / sparse array / tuple, optional (device or host) - Either one of a tuple (indices, distances) of - arrays of shape (n_samples, n_neighbors), a pairwise distances - dense array of shape (n_samples, n_samples) or a KNN graph - sparse array (preferably CSR/COO). This feature allows - the precomputation of the KNN outside of TSNE - and also allows the use of a custom distance function. This function - should match the metric used to train the TSNE embeedings. - Takes precedence over the precomputed_knn parameter. + knn_graph : tuple[array, array], sparse-matrix, array, optional + This feature allows the precomputation of the KNN outside of TSNE. + + This may take any of the valid forms accepted by the + ``precomputed_knn`` parameter to ``TSNE``, and takes precedence + over it. See the ``TSNE`` docstring on ``precomputed_knn`` for more + information. """ cdef int n_samples, n_features cdef uintptr_t X_ptr = 0 @@ -610,20 +620,13 @@ class TSNE(InteropMixin, if knn_graph is None: knn_graph = self.precomputed_knn if knn_graph is not None: - knn_indices, knn_dists = extract_knn_graph(knn_graph, params.n_neighbors) - - knn_dists_cp = knn_dists.to_output("cupy") - - if sparse_fit: - # Sparse fitting requires the indices to be int32 - knn_indices_cp = cupy.asarray( - knn_indices.to_output("cupy"), dtype=np.int32 - ) - else: - knn_indices_cp = knn_indices.to_output("cupy") - - knn_dists_ptr = knn_dists_cp.data.ptr - knn_indices_ptr = knn_indices_cp.data.ptr + knn_indices, knn_dists = extract_knn_graph( + knn_graph, + params.n_neighbors, + indices_dtype="int32" if sparse_fit else "int64", + ) + knn_dists_ptr = knn_dists.data.ptr + knn_indices_ptr = knn_indices.data.ptr # Allocate output array embedding = cupy.zeros( diff --git a/python/cuml/cuml/manifold/umap/umap.pyx b/python/cuml/cuml/manifold/umap/umap.pyx index 25c2fa72a9..f90673b697 100644 --- a/python/cuml/cuml/manifold/umap/umap.pyx +++ b/python/cuml/cuml/manifold/umap/umap.pyx @@ -6,8 +6,6 @@ import ctypes import warnings from collections import deque -from cuda.bindings.cyruntime cimport cudaStream_t - import cupy as cp import cupyx.scipy.sparse import joblib @@ -18,7 +16,6 @@ import scipy.spatial from cuml.common.array_descriptor import CumlArrayDescriptor from cuml.common.doc_utils import generate_docstring from cuml.common.sparse_utils import is_sparse -from cuml.common.sparsefuncs import extract_knn_graph from cuml.internals import logger, reflect from cuml.internals.array import CumlArray from cuml.internals.array_sparse import SparseCumlArray @@ -39,7 +36,9 @@ from cuml.internals.validation import ( check_random_seed, check_y, ) +from cuml.manifold.utils import extract_knn_graph +from cuda.bindings.cyruntime cimport cudaStream_t from libc.stdint cimport int64_t, uintptr_t from libcpp cimport bool from libcpp.memory cimport unique_ptr @@ -790,16 +789,27 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): feature is made optional in the GPU version due to the significant overhead in copying memory to the host for computing the hash. - precomputed_knn : array / sparse array / tuple, optional (device or host) - Either one of a tuple (indices, distances) of - arrays of shape (n_samples, n_neighbors), a pairwise distances - dense array of shape (n_samples, n_samples) or a KNN graph - sparse array (preferably CSR/COO). This feature allows - the precomputation of the KNN outside of UMAP - and also allows the use of a custom distance function. This function - should match the metric used to train the UMAP embeddings. For most efficient - memory usage, the precomputed knn graph should be CPU-accessible arrays - such as numpy arrays. + precomputed_knn : tuple[array, array], sparse-matrix, array, optional + This feature allows the precomputation of the KNN outside of UMAP. + Options are: + + - A tuple (indices, distances) of dense arrays of shape (n_samples, + n_neighbors), where n_neighbors is >= the ``n_neighbors`` parameter. + Self references should be included (i.e. the first column of + `indices` should be [0, 1, ...], denotating that the nearest neighbor + to each row is itself). This is the most efficient representation. + Note that providing on CPU may result in lower peak GPU memory usage. + + - A sparse matrix KNN graph, as may be output by + ``cuml.neighbors.kneighbors_graph`` with ``mode="distance"`` and + ``include_self=True``. The ``n_neighbors`` used to calculate the + graph must be >= the ``n_neighbors`` parameter. + + - A pairwise distances dense array of shape (n_samples, n_samples). + + In all cases the KNN should be computed using the same ``metric`` as + provided to ``UMAP``. + random_state : int, RandomState instance or None, optional (default=None) Seed used by the random number generator for embedding initialization and optimizer sampling. Setting a random_state enables reproducible @@ -1084,10 +1094,10 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): input_hash = _joblib_hash(raw_data) if (knn_dists := getattr(self, "_knn_dists", None)) is not None: - knn_dists = to_cpu(knn_dists) + knn_dists = cp.asnumpy(knn_dists) if (knn_indices := getattr(self, "_knn_indices", None)) is not None: - knn_indices = to_cpu(knn_indices) + knn_indices = cp.asnumpy(knn_indices) attrs = { "embedding_": to_cpu(self.embedding_), @@ -1205,17 +1215,13 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): Parameters ---------- - knn_graph : array / sparse array / tuple, optional (device or host) - Either one of a tuple (indices, distances) of - arrays of shape (n_samples, n_neighbors), a pairwise distances - dense array of shape (n_samples, n_samples) or a KNN graph - sparse array (preferably CSR/COO). This feature allows - the precomputation of the KNN outside of UMAP - and also allows the use of a custom distance function. This function - should match the metric used to train the UMAP embeddings. - Takes precedence over the precomputed_knn parameter. For most efficient - memory usage, the precomputed knn graph should be CPU-accessible arrays - such as numpy arrays. + knn_graph: tuple[array, array], sparse-matrix, array, optional + This feature allows the precomputation of the KNN outside of UMAP. + + This may take any of the valid forms accepted by the + ``precomputed_knn`` parameter to ``UMAP``, and takes precedence + over it. See the ``UMAP`` docstring on ``precomputed_knn`` for more + information. """ # Normalize X as cheaply as possible to minimize copies and work X, index = check_inputs( @@ -1302,19 +1308,15 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): knn_indices, knn_dists = extract_knn_graph( (knn_graph if knn_graph is not None else self.precomputed_knn), self._n_neighbors, - mem_type=False, # mirrors the input graph mem type + mem_type=None, # mirrors the input graph mem type + indices_dtype=("int32" if X_is_sparse else "int64"), ) - knn_dists_cp = knn_dists.to_output("cupy") - if X_is_sparse: - knn_indices_cp = cp.asarray( - knn_indices.to_output("cupy"), dtype=np.int32 - ) - # Drop the int64 original and keep only the int32 copy used by the kernel. - knn_indices = CumlArray(data=knn_indices_cp) + if isinstance(knn_indices, cp.ndarray): + knn_indices_ptr = knn_indices.data.ptr + knn_dists_ptr = knn_dists.data.ptr else: - knn_indices_cp = knn_indices.to_output("cupy") - knn_indices_ptr = knn_indices_cp.data.ptr - knn_dists_ptr = knn_dists_cp.data.ptr + knn_indices_ptr = knn_indices.ctypes.data + knn_dists_ptr = knn_dists.ctypes.data else: knn_indices = knn_dists = None @@ -1451,17 +1453,13 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): Parameters ---------- - knn_graph : array / sparse array / tuple, optional (device or host) - Either one of a tuple (indices, distances) of - arrays of shape (n_samples, n_neighbors), a pairwise distances - dense array of shape (n_samples, n_samples) or a KNN graph - sparse array (preferably CSR/COO). This feature allows - the precomputation of the KNN outside of UMAP - and also allows the use of a custom distance function. This function - should match the metric used to train the UMAP embeddings. - Takes precedence over the precomputed_knn parameter. For most efficient - memory usage, the precomputed knn graph should be CPU-accessible arrays - such as numpy arrays. + knn_graph : tuple[array, array], sparse-matrix, array, optional + This feature allows the precomputation of the KNN outside of UMAP. + + This may take any of the valid forms accepted by the + ``precomputed_knn`` parameter to ``UMAP``, and takes precedence + over it. See the ``UMAP`` docstring on ``precomputed_knn`` for more + information. """ self.fit(X, y, convert_dtype=convert_dtype, knn_graph=knn_graph) return self.embedding_ diff --git a/python/cuml/cuml/manifold/utils.py b/python/cuml/cuml/manifold/utils.py new file mode 100644 index 0000000000..19d62ebf79 --- /dev/null +++ b/python/cuml/cuml/manifold/utils.py @@ -0,0 +1,279 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# +import math + +import cupy as cp +import numpy as np + +from cuml.common.kernel_utils import cuda_kernel_factory +from cuml.common.sparse_utils import is_sparse +from cuml.internals.validation import check_array + +_is_axis_1_sorted_kernel = cp.RawKernel( + """ + extern "C" __global__ + void is_axis_1_sorted(const float* arr, int n_rows, int n_cols, int *sorted) { + int row = blockDim.x * blockIdx.x + threadIdx.x; + + if (row >= n_rows) return; + + int start = row * n_cols; + int end = start + n_cols - 1; + for (int i = start; i < end; i++) { + if (arr[i] > arr[i + 1]) { + *sorted = 0; + return; + } + } + } + """, + "is_axis_1_sorted", +) + + +def _is_axis_1_sorted(X): + """Checks if axis 1 (every row) is sorted in X.""" + # Safety checks + assert X.flags.c_contiguous + assert X.dtype == "float32" + assert X.ndim == 2 + # XXX: ensure on device just for this routine - zero copy unless a + # host->device transfer needed. + X = cp.asarray(X, order="C") + is_sorted = cp.ones(1, dtype="int32") + _is_axis_1_sorted_kernel( + (math.ceil(X.shape[0] / 32),), + (32,), + (X, X.shape[0], X.shape[1], is_sorted), + ) + return bool(is_sorted.item()) + + +def _has_self_references(indices): + """Checks if indices has self references in column 0.""" + # Safety checks + assert indices.flags.c_contiguous + assert indices.ndim == 2 + n_rows, n_cols = indices.shape + if isinstance(indices, cp.ndarray): + out = cp.ones(1, dtype="int32") + kernel = cuda_kernel_factory( + """ + (const {0}* arr, int n_rows, int n_cols, int *out) { + int row = blockDim.x * blockIdx.x + threadIdx.x; + + if (row >= n_rows) return; + if (arr[row * n_cols] != row) { + *out = 0; + } + return; + } + """, + (indices.dtype,), + "has_self_references", + ) + kernel( + (math.ceil(n_rows / 32),), (32,), (indices, n_rows, n_cols, out) + ) + return bool(out.item()) + return (indices[:, 0] == np.arange(n_rows, dtype=indices.dtype)).all() + + +def extract_knn_graph( + knn_info, + n_neighbors, + mem_type="device", + indices_dtype="int64", +): + """Extract the KNN graph indices and distances. + + Parameters + ---------- + knn_info : array, sparse-matrix, or tuple[array, array] + - Tuple (indices, distances) of arrays of shape (n_samples, + n_neighbors). Should contain self references (i.e. the closest sample + to a row is the row itself). + - Pairwise distances dense array of shape (n_samples, n_samples). + - KNN graph sparse array. This is most efficient if the graph is in CSR + format and contains 0 entries in `data` for all diagonal elements. + n_neighbors: int + Number of nearest neighbors + mem_type : {"device", "host", None}, default="device" + The desired output memory type. + indices_dtype : dtype, default='int64' + The dtype to use for the output indices. + + Returns + ------- + indices : cupy.ndarray or numpy.ndarray + The KNN indices, shape=n_samples * n_neighbors, dtype=indices_dtype. + distances : cupy.ndarray or numpy.ndarray + The KNN distances, shape=n_samples * n_neighbors, dtype=float32. + """ + # The initial mem_type to coerce to. When possible we only coerce to device + # if the output is known to be device, otherwise we leave as is until the + # final coercion. + mem_type_init = "device" if mem_type == "device" else None + if isinstance(knn_info, tuple): + # (indices, distances), each with shape=(n_samples, orig_n_neighbors) + indices, distances = knn_info + indices = check_array( + indices, dtype=indices_dtype, order="C", mem_type=mem_type_init + ) + distances = check_array( + distances, dtype="float32", order="C", mem_type=mem_type_init + ) + if not indices.shape == distances.shape: + raise ValueError( + f"Expected indices and distances to have shape=(n_samples, " + f"n_neighbors), got indices.shape={indices.shape}, " + f"distances.shape={distances.shape}" + ) + if not _has_self_references(indices): + raise ValueError( + "Expected indices and distances to include self references (i.e. " + "the closest sample to each row is itself). If using " + "`NearestNeighbors.kneighbors` to precompute the KNN, pass in " + "the training data to both `NearestNeighbors.fit` and " + "`NearestNeighbors.kneighbors`." + ) + elif is_sparse(knn_info): + # Sparse KNN graph + # - shape=(n_samples, n_samples) + # - nnz=n_samples * orig_n_neighbors + if not (knn_info.ndim == 2 and knn_info.shape[0] == knn_info.shape[1]): + raise ValueError( + f"Expected a sparse array of shape=(n_samples, n_samples), " + f"got shape={knn_info.shape}" + ) + + # Coerce to CSR. If the input was already CSR this is zero-copy and + # avoids reordering indices (leaving `.data` in the initial order). + # This ensures the case of passing a direct `kneighbors_graph` output + # can be done zero-copy. + knn_info = knn_info.tocsr() + + # XXX: Here we try to do some rudimentary inference and validation of + # the original K on the provided KNN graph. Since sparse inputs may + # drop 0 entries, some diagonal entries may be _close_ to zero but not + # exactly zero, and the graph may or may not be canonicalized this is + # only a best-effort. Some invalid graphs may sneak through here, but + # we try to catch common errors and issues and raise them. + n_samples = knn_info.shape[0] + if (remainder := knn_info.nnz % n_samples) != 0: + # This might still be valid if _some_ diagonal distances are + # _close_ to zero but not actually zero, _and_ the actually zero + # elements were dropped. Before doing a copy, check if dropping + # the diagonal will still result in an invalid graph. + if (knn_info.diagonal() != 0).sum() != remainder: + raise ValueError( + f"Precomputed KNN graph has {knn_info.nnz} nonzero elements which " + f"is not evenly divisible by {n_samples} samples." + ) + knn_info = knn_info.copy() + knn_info.setdiag(knn_info.dtype.type(0)) + knn_info.eliminate_zeros() + orig_n_neighbors = knn_info.nnz // n_samples + + indices = check_array( + knn_info.indices.reshape((n_samples, orig_n_neighbors)), + dtype=indices_dtype, + order="C", + mem_type=mem_type_init, + ) + distances = check_array( + knn_info.data.reshape((n_samples, orig_n_neighbors)), + dtype="float32", + order="C", + mem_type=mem_type_init, + ) + xp = cp if isinstance(distances, cp.ndarray) else np + # Reorder by distances if not already sorted. This is necessary for KNN + # graph inputs, since a canonical sparse matrix will not have the data + # sorted as we need it. Both sklearn and cuml's `kneighbors_graph` + # returns a matrix with `.data` sorted as required (not canonicalized), + # so we optimistically check for sortedness before doing the sorting. + if not _is_axis_1_sorted(distances): + new_order = distances.argsort() + all_rows = xp.arange(distances.shape[0])[:, None] + indices = indices[all_rows, new_order] + distances = distances[all_rows, new_order] + del new_order + + if not _has_self_references(indices): + # No self references present in the data. The KNN graph was either + # computed without self references, or 0 elements were dropped + # during a canonicalization step at some point. + temp = xp.zeros( + (n_samples, orig_n_neighbors + 1), + dtype=distances.dtype, + order="C", + ) + temp[:, 1:] = distances + distances = temp + temp = xp.empty( + (n_samples, orig_n_neighbors + 1), + dtype=indices.dtype, + order="C", + ) + temp[:, 0] = xp.arange(n_samples, dtype=indices.dtype) + temp[:, 1:] = indices + indices = temp + orig_n_neighbors += 1 + else: + # Dense pairwise distance matrix, shape=(n_samples, n_samples) + knn_info = check_array( + knn_info, dtype="float32", mem_type=mem_type_init + ) + if knn_info.shape[0] != knn_info.shape[1]: + raise ValueError( + f"Expected a dense array of shape=(n_samples, n_samples), " + f"got shape={knn_info.shape}" + ) + n_samples = knn_info.shape[0] + if n_samples < n_neighbors: + raise ValueError( + f"Precomputed KNN data requires n_samples >= n_neighbors. " + f"Got {n_neighbors=}, {n_samples=}" + ) + + # Convert pairwise distance matrix to KNN graph + xp = cp if isinstance(knn_info, cp.ndarray) else np + # Partition indices to select the nearest `n_neighbors` + indices = xp.argpartition(knn_info, n_neighbors - 1, axis=1) + indices = indices[:, :n_neighbors] + # Reorder and subset indices and distances appropriately + all_rows = xp.arange(n_samples)[:, None] + indices = indices[all_rows, xp.argsort(knn_info[all_rows, indices])] + distances = knn_info[all_rows, indices] + + # Validate shape and n_neighbors + if indices.shape[1] < n_neighbors: + raise ValueError( + f"Precomputed KNN data has {indices.shape[1]} neighbors per " + f"sample, but {n_neighbors=} was specified. Please provide KNN data " + f"with at least {n_neighbors} neighbors per sample." + ) + + # Trim arrays to n_neighbors if necessary + if indices.shape[1] > n_neighbors: + indices = indices[:, :n_neighbors] + distances = distances[:, :n_neighbors] + + # Reshape and coerce to proper dtype and mem_type. + indices = check_array( + indices.reshape(-1), + dtype=indices_dtype, + mem_type=mem_type, + ensure_2d=False, + ) + distances = check_array( + distances.reshape(-1), + dtype="float32", + mem_type=("device" if isinstance(indices, cp.ndarray) else "host"), + ensure_2d=False, + ensure_all_finite=False, + ) + return indices, distances diff --git a/python/cuml/tests/test_umap.py b/python/cuml/tests/test_umap.py index 0cb6b04a44..4f3746f7a8 100644 --- a/python/cuml/tests/test_umap.py +++ b/python/cuml/tests/test_umap.py @@ -6,12 +6,12 @@ from importlib.metadata import version as package_version import cupy as cp -import cupyx +import cupyx.scipy.sparse as cp_sp import joblib import numba import numpy as np import pytest -import scipy.sparse as scipy_sparse +import scipy.sparse as sp import umap from packaging.version import Version from sklearn import datasets @@ -19,11 +19,12 @@ from sklearn.datasets import make_blobs, make_moons from sklearn.manifold import trustworthiness from sklearn.metrics import adjusted_rand_score -from sklearn.neighbors import KDTree, NearestNeighbors +from sklearn.neighbors import KDTree, NearestNeighbors, kneighbors_graph import cuml from cuml.internals import GraphBasedDimRedCallback from cuml.manifold.umap import UMAP as cuUMAP +from cuml.manifold.utils import extract_knn_graph from cuml.metrics import pairwise_distances from cuml.testing.utils import ( array_equal, @@ -178,14 +179,9 @@ def test_umap_transform_on_digits_sparse( [True, False], 1797, replace=True, p=[0.75, 0.25] ) - if input_type == "cupy": - sp_prefix = cupyx.scipy.sparse - else: - sp_prefix = scipy_sparse + sparse = cp_sp if input_type == "cupy" else sp - data = sp_prefix.csr_matrix( - scipy_sparse.csr_matrix(digits.data[digits_selection]) - ) + data = sparse.csr_matrix(sp.csr_matrix(digits.data[digits_selection])) fitter = cuUMAP( n_neighbors=15, @@ -196,9 +192,7 @@ def test_umap_transform_on_digits_sparse( target_metric=target_metric, ) - new_data = sp_prefix.csr_matrix( - scipy_sparse.csr_matrix(digits.data[~digits_selection]) - ) + new_data = sparse.csr_matrix(sp.csr_matrix(digits.data[~digits_selection])) if xform_method == "fit": fitter.fit(data, convert_dtype=True) @@ -585,6 +579,199 @@ def test_fit_fewer_rows_than_n_neighbors(): assert model._n_neighbors == 10 +@pytest.mark.parametrize("in_mem_type", ["device", "host"]) +@pytest.mark.parametrize("in_dtype", ["float32", "float64"]) +@pytest.mark.parametrize("indices_dtype", ["float32", "float64"]) +@pytest.mark.parametrize("k", [10, 20]) +@pytest.mark.parametrize( + "kind, transform", + [ + ("csr", None), + ("csc", None), + ("coo", None), + ("csr", "canonical"), + ("csc", "canonical"), + ("coo", "canonical"), + ("csr", "nonzero"), + ("csc", "nonzero"), + ("coo", "nonzero"), + ("pairwise", None), + ("tuple", None), + ], +) +def test_extract_knn_graph( + in_mem_type, in_dtype, indices_dtype, k, kind, transform +): + X, _ = datasets.make_blobs(30, 10, centers=5, random_state=42) + if kind == "pairwise": + knn_info = pairwise_distances(X).astype(in_dtype) + if in_mem_type == "device": + knn_info = cp.asarray(knn_info) + elif kind == "tuple": + nn = NearestNeighbors(n_neighbors=20).fit(X) + distances, indices = nn.kneighbors(X, return_distance=True) + if in_mem_type == "device": + indices = cp.asarray(indices) + distances = cp.asarray(distances) + knn_info = (indices, distances.astype(in_dtype)) + else: + knn_info = kneighbors_graph(X, 20, include_self=True, mode="distance") + if in_mem_type == "device": + knn_info = cp_sp.csr_matrix(knn_info) + knn_info = knn_info.asformat(kind).astype(in_dtype) + if transform is not None: + # Canonicalize input matrix + if hasattr(knn_info, "sort_indices"): + knn_info.sort_indices() + if transform == "nonzero": + knn_info.eliminate_zeros() + + sol = kneighbors_graph(X, k, include_self=True, mode="distance") + + # mem_type=None doesn't coerce mem_type + indices, distances = extract_knn_graph( + knn_info, k, indices_dtype=indices_dtype, mem_type=None + ) + assert indices.dtype == indices_dtype + assert distances.dtype == "float32" + xp = cp if in_mem_type == "device" else np + assert isinstance(indices, xp.ndarray) + assert isinstance(distances, xp.ndarray) + np.testing.assert_allclose(cp.asnumpy(distances), sol.data, atol=1e-5) + np.testing.assert_array_equal(cp.asnumpy(indices), sol.indices) + + indices, distances = extract_knn_graph( + knn_info, k, indices_dtype=indices_dtype, mem_type="host" + ) + assert indices.dtype == indices_dtype + assert distances.dtype == "float32" + np.testing.assert_allclose(distances, sol.data, atol=1e-5) + np.testing.assert_array_equal(indices, sol.indices) + + indices, distances = extract_knn_graph( + knn_info, k, indices_dtype=indices_dtype, mem_type="device" + ) + assert indices.dtype == indices_dtype + assert distances.dtype == "float32" + np.testing.assert_allclose(distances.get(), sol.data, atol=1e-5) + np.testing.assert_array_equal(indices.get(), sol.indices) + + +def test_extract_knn_graph_errors(): + X, _ = datasets.make_blobs(30, 10, centers=5, random_state=42) + + # tuple: invalid shape + knn_info = (np.ones((3, 4)), np.ones((4, 3))) + with pytest.raises( + ValueError, match="Expected indices and distances to have shape" + ): + extract_knn_graph(knn_info, 10) + + # tuple: n_neighbors > original_n_neighbors + knn_info = tuple( + reversed( + NearestNeighbors(n_neighbors=20) + .fit(X) + .kneighbors(X, return_distance=True) + ) + ) + with pytest.raises( + ValueError, + match=( + "Precomputed KNN data has 20 neighbors per sample, but " + "n_neighbors=25 was specified" + ), + ): + extract_knn_graph(knn_info, 25) + + # tuple: expected self references + knn_info = tuple( + reversed( + NearestNeighbors(n_neighbors=20) + .fit(X) + .kneighbors(return_distance=True) + ) + ) + with pytest.raises( + ValueError, + match="Expected indices and distances to include self references", + ): + extract_knn_graph(knn_info, 20) + + # graph: invalid shape + knn_info = sp.random(10, 25, density=0.5) + with pytest.raises(ValueError, match="Expected a sparse array of shape"): + extract_knn_graph(knn_info, 10) + + # graph: invalid nnz + knn_info = sp.csr_matrix( + ( + np.array([0.5, 0.5, 0.5, 0.5]), + (np.array([0, 0, 1, 1]), np.array([1, 2, 0, 2])), + ), + shape=(3, 3), + ) + with pytest.raises( + ValueError, match="Precomputed KNN graph has 4 nonzero elements" + ): + extract_knn_graph(knn_info, 3) + + # graph: n_neighbors > original_n_neighbors + knn_info = kneighbors_graph(X, 20, include_self=True, mode="distance") + with pytest.raises( + ValueError, + match=( + "Precomputed KNN data has 20 neighbors per sample, but " + "n_neighbors=25 was specified" + ), + ): + extract_knn_graph(knn_info, 25) + + # pairwise: invalid shape + knn_info = np.ones((5, 6)) + with pytest.raises(ValueError, match="Expected a dense array of shape"): + extract_knn_graph(knn_info, 5) + + # pairwise: n_neighbors > n_samples + knn_info = np.ones((5, 5)) + with pytest.raises( + ValueError, + match="Precomputed KNN data requires n_samples >= n_neighbors", + ): + extract_knn_graph(knn_info, 10) + + +@pytest.mark.parametrize( + "kind, mem_type", + [ + ("cupy", None), + ("cupy", "device"), + ("numpy", None), + ("numpy", "host"), + ], +) +def test_extract_knn_graph_from_kneighbors_graph_zero_copy(kind, mem_type): + """When converting from the output of `kneighbors_graph`, no copy should be + needed if not changing mem_type, as no sorting is required""" + X, _ = datasets.make_blobs(30, 10, centers=5, random_state=42) + with cuml.using_output_type(kind): + graph = cuml.neighbors.kneighbors_graph( + X, 20, include_self=True, mode="distance" + ) + + def assert_same_memory(x, y): + assert type(x) is type(y) + x_ptr = x.data.ptr if isinstance(x, cp.ndarray) else x.ctypes.data + y_ptr = y.data.ptr if isinstance(y, cp.ndarray) else y.ctypes.data + assert x_ptr == y_ptr + + indices, distances = extract_knn_graph( + graph, 20, mem_type=mem_type, indices_dtype=graph.indices.dtype + ) + assert_same_memory(indices, graph.indices) + assert_same_memory(distances, graph.data) + + @pytest.mark.parametrize("n_neighbors", [5, 15]) @pytest.mark.parametrize("build_algo", ["brute_force_knn", "nn_descent"]) @pytest.mark.parametrize("data_on_gpu", [True, False]) @@ -682,7 +869,7 @@ def test_umap_precomputed_knn(precomputed_type, sparse_input, build_algo): [0.0, 1.0], p=[0.1, 0.9], size=data.shape ) data = np.multiply(data, sparsification) - data = scipy_sparse.csr_matrix(data) + data = sp.csr_matrix(data) n_neighbors = 8 @@ -866,7 +1053,7 @@ def test_umap_distance_metrics_fit_transform_trust_on_sparse_input( if metric == "jaccard": data = data >= 0 - new_data = scipy_sparse.csr_matrix(data[~data_selection]) + new_data = sp.csr_matrix(data[~data_selection]) if umap_learn_supported: umap_model = umap.UMAP( @@ -1199,7 +1386,9 @@ def test_umap_precomputed_knn_insufficient_neighbors(precomputed_type): random_state=42, init="random", ) - with pytest.raises(ValueError, match=".*fewer neighbors.*"): + with pytest.raises( + ValueError, match=f".*at least {k_requested} neighbors.*" + ): model.fit(data) @@ -1395,9 +1584,7 @@ def to_np(arr): def test_inverse_transform_sparse_error(): """Test that inverse_transform raises error for sparse input data.""" # Create sparse data - X_sparse = scipy_sparse.random( - 100, 20, density=0.3, format="csr", random_state=42 - ) + X_sparse = sp.random(100, 20, density=0.3, format="csr", random_state=42) X_sparse = X_sparse.astype(np.float32) # Fit UMAP on sparse data From 8cc749d1afa46fb56fe040ab79aefee1964e2954 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Tue, 23 Jun 2026 12:39:29 -0500 Subject: [PATCH 2/3] Respond to feedback --- python/cuml/cuml/manifold/utils.py | 162 ++++++++++++++++------------- python/cuml/tests/test_umap.py | 26 +++++ 2 files changed, 116 insertions(+), 72 deletions(-) diff --git a/python/cuml/cuml/manifold/utils.py b/python/cuml/cuml/manifold/utils.py index 19d62ebf79..d33d8d9bbb 100644 --- a/python/cuml/cuml/manifold/utils.py +++ b/python/cuml/cuml/manifold/utils.py @@ -11,6 +11,35 @@ from cuml.common.sparse_utils import is_sparse from cuml.internals.validation import check_array + +def _check_indices_per_row(indptr, n_neighbors): + """Check if indptr indicates n_neighbors per row""" + if isinstance(indptr, cp.ndarray): + out = cp.ones(1, dtype="int32") + kernel = cuda_kernel_factory( + """ + (const {0}* arr, int n_rows, int n_neighbors, int *out) { + int row = blockDim.x * blockIdx.x + threadIdx.x; + + if (row + 1 >= n_rows) return; + if (arr[row + 1] - arr[row] != n_neighbors) { + *out = 0; + } + return; + } + """, + (indptr.dtype,), + "has_n_neighbors_per_row", + ) + kernel( + (math.ceil(indptr.shape[0] / 32),), + (32,), + (indptr, indptr.shape[0], n_neighbors, out), + ) + return bool(out.item()) + return (np.diff(indptr) == n_neighbors).all() + + _is_axis_1_sorted_kernel = cp.RawKernel( """ extern "C" __global__ @@ -33,29 +62,26 @@ ) -def _is_axis_1_sorted(X): - """Checks if axis 1 (every row) is sorted in X.""" - # Safety checks - assert X.flags.c_contiguous - assert X.dtype == "float32" - assert X.ndim == 2 +def _check_distances_sorted(dist, orig_n_neighbors): + """Checks if axis 1 (every row) is sorted appropriately in dist.""" + # Safety check + assert dist.dtype == "float32" # XXX: ensure on device just for this routine - zero copy unless a # host->device transfer needed. - X = cp.asarray(X, order="C") + dist = cp.asarray(dist, order="C") is_sorted = cp.ones(1, dtype="int32") _is_axis_1_sorted_kernel( - (math.ceil(X.shape[0] / 32),), + (math.ceil(dist.shape[0] / 32),), (32,), - (X, X.shape[0], X.shape[1], is_sorted), + (dist, dist.shape[0] // orig_n_neighbors, orig_n_neighbors, is_sorted), ) return bool(is_sorted.item()) -def _has_self_references(indices): +def _check_self_references(indices, orig_n_neighbors=None): """Checks if indices has self references in column 0.""" - # Safety checks - assert indices.flags.c_contiguous - assert indices.ndim == 2 + if orig_n_neighbors is not None: + indices = indices.reshape((-1, orig_n_neighbors)) n_rows, n_cols = indices.shape if isinstance(indices, cp.ndarray): out = cp.ones(1, dtype="int32") @@ -131,7 +157,7 @@ def extract_knn_graph( f"n_neighbors), got indices.shape={indices.shape}, " f"distances.shape={distances.shape}" ) - if not _has_self_references(indices): + if not _check_self_references(indices): raise ValueError( "Expected indices and distances to include self references (i.e. " "the closest sample to each row is itself). If using " @@ -153,75 +179,67 @@ def extract_knn_graph( # avoids reordering indices (leaving `.data` in the initial order). # This ensures the case of passing a direct `kneighbors_graph` output # can be done zero-copy. - knn_info = knn_info.tocsr() + knn_info = check_array( + knn_info, + accept_sparse=["csr"], + accept_large_sparse=(indices_dtype == "int64"), + mem_type=mem_type_init, + dtype="float32", + ) + xp = cp if isinstance(knn_info.data, cp.ndarray) else np - # XXX: Here we try to do some rudimentary inference and validation of - # the original K on the provided KNN graph. Since sparse inputs may - # drop 0 entries, some diagonal entries may be _close_ to zero but not - # exactly zero, and the graph may or may not be canonicalized this is - # only a best-effort. Some invalid graphs may sneak through here, but - # we try to catch common errors and issues and raise them. n_samples = knn_info.shape[0] - if (remainder := knn_info.nnz % n_samples) != 0: - # This might still be valid if _some_ diagonal distances are - # _close_ to zero but not actually zero, _and_ the actually zero - # elements were dropped. Before doing a copy, check if dropping - # the diagonal will still result in an invalid graph. - if (knn_info.diagonal() != 0).sum() != remainder: + + orig_n_neighbors, remainder = divmod(knn_info.nnz, n_samples) + if ( + remainder == 0 + and _check_indices_per_row(knn_info.indptr, orig_n_neighbors) + and _check_distances_sorted(knn_info.data, orig_n_neighbors) + and _check_self_references(knn_info.indices, orig_n_neighbors) + ): + # The input graph is usable as is with no copies needed. + distances = knn_info.data.reshape((n_samples, orig_n_neighbors)) + indices = knn_info.indices.reshape((n_samples, orig_n_neighbors)) + else: + # Graph `data` and `indices` aren't in the correct format, we + # need to copy and massage the data. + knn_info = knn_info.copy() + + # Set an explicit value for all diagonal elements, ensuring self + # references are present. We use -1 to ensure that self references + # sort earlier than any other 0 distance elements. + knn_info.setdiag(knn_info.dtype.type(-1)) + + # Recompute and validate orig_n_neighbors. An error here indicates + # either an invalid KNN graph or one where some samples had 0 + # distance to each other _and_ the 0 entries were dropped. We cannot + # recover that information and have to error. + orig_n_neighbors, remainder = divmod(knn_info.nnz, n_samples) + if not ( + remainder == 0 + and _check_indices_per_row(knn_info.indptr, orig_n_neighbors) + ): raise ValueError( - f"Precomputed KNN graph has {knn_info.nnz} nonzero elements which " - f"is not evenly divisible by {n_samples} samples." + f"Precomputed KNN graph has {knn_info.nnz - n_samples} " + f"nonzero elements which is not evenly divisible by " + f"{n_samples} samples. The graph may be malformed, or may " + f"have contained 0-distance samples that were removed " + f"during sparse matrix canonicalization." ) - knn_info = knn_info.copy() - knn_info.setdiag(knn_info.dtype.type(0)) - knn_info.eliminate_zeros() - orig_n_neighbors = knn_info.nnz // n_samples - indices = check_array( - knn_info.indices.reshape((n_samples, orig_n_neighbors)), - dtype=indices_dtype, - order="C", - mem_type=mem_type_init, - ) - distances = check_array( - knn_info.data.reshape((n_samples, orig_n_neighbors)), - dtype="float32", - order="C", - mem_type=mem_type_init, - ) - xp = cp if isinstance(distances, cp.ndarray) else np - # Reorder by distances if not already sorted. This is necessary for KNN - # graph inputs, since a canonical sparse matrix will not have the data - # sorted as we need it. Both sklearn and cuml's `kneighbors_graph` - # returns a matrix with `.data` sorted as required (not canonicalized), - # so we optimistically check for sortedness before doing the sorting. - if not _is_axis_1_sorted(distances): + # Extract distances and indices into individual arrays + distances = knn_info.data.reshape((n_samples, orig_n_neighbors)) + indices = knn_info.indices.reshape((n_samples, orig_n_neighbors)) + + # Sort each row by distance. new_order = distances.argsort() all_rows = xp.arange(distances.shape[0])[:, None] indices = indices[all_rows, new_order] distances = distances[all_rows, new_order] del new_order - if not _has_self_references(indices): - # No self references present in the data. The KNN graph was either - # computed without self references, or 0 elements were dropped - # during a canonicalization step at some point. - temp = xp.zeros( - (n_samples, orig_n_neighbors + 1), - dtype=distances.dtype, - order="C", - ) - temp[:, 1:] = distances - distances = temp - temp = xp.empty( - (n_samples, orig_n_neighbors + 1), - dtype=indices.dtype, - order="C", - ) - temp[:, 0] = xp.arange(n_samples, dtype=indices.dtype) - temp[:, 1:] = indices - indices = temp - orig_n_neighbors += 1 + # Finally swap -1 distances for self references back to 0 + distances[:, 0] = 0 else: # Dense pairwise distance matrix, shape=(n_samples, n_samples) knn_info = check_array( diff --git a/python/cuml/tests/test_umap.py b/python/cuml/tests/test_umap.py index 4f3746f7a8..def7fb8968 100644 --- a/python/cuml/tests/test_umap.py +++ b/python/cuml/tests/test_umap.py @@ -741,6 +741,32 @@ def test_extract_knn_graph_errors(): extract_knn_graph(knn_info, 10) +@pytest.mark.parametrize( + "kind, mem_type", + [ + ("cupy", None), + ("cupy", "device"), + ("numpy", None), + ("numpy", "host"), + ], +) +def test_extract_knn_graph_zero_distance_rows(kind, mem_type): + X = np.array([[0, 1, 2], [3, 1, 3], [5, 1, 2], [0, 1, 2]], dtype="float32") + nn = cuml.neighbors.NearestNeighbors().fit(X) + with cuml.using_output_type(kind): + graph = nn.kneighbors_graph(X, 3, mode="distance") + graph.sort_indices() + + # Hardcoded indices to ease testing - it's tricky to get `kneighbors` to + # return self references first otherwise. + sol_indices = np.array([0, 3, 1, 1, 2, 3, 2, 1, 3, 3, 0, 1]) + sol_distances = nn.kneighbors(X, 3)[0].flatten() + + indices, distances = extract_knn_graph(graph, 3, mem_type=mem_type) + np.testing.assert_allclose(cp.asnumpy(distances), sol_distances, atol=1e-5) + np.testing.assert_array_equal(cp.asnumpy(indices), sol_indices) + + @pytest.mark.parametrize( "kind, mem_type", [ From c7fbee9d18dedfd94d3375ddc9a76685bb767005 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Tue, 23 Jun 2026 15:52:42 -0500 Subject: [PATCH 3/3] Add check for correct `n_samples`, and fix TSNE test --- python/cuml/cuml/manifold/t_sne.pyx | 1 + python/cuml/cuml/manifold/umap/umap.pyx | 1 + python/cuml/cuml/manifold/utils.py | 29 +++++++++--------- python/cuml/tests/test_tsne.py | 39 ++++--------------------- python/cuml/tests/test_umap.py | 39 ++++++++++++++++--------- 5 files changed, 49 insertions(+), 60 deletions(-) diff --git a/python/cuml/cuml/manifold/t_sne.pyx b/python/cuml/cuml/manifold/t_sne.pyx index e50af9bf85..1e20374f4a 100644 --- a/python/cuml/cuml/manifold/t_sne.pyx +++ b/python/cuml/cuml/manifold/t_sne.pyx @@ -622,6 +622,7 @@ class TSNE(InteropMixin, if knn_graph is not None: knn_indices, knn_dists = extract_knn_graph( knn_graph, + n_samples, params.n_neighbors, indices_dtype="int32" if sparse_fit else "int64", ) diff --git a/python/cuml/cuml/manifold/umap/umap.pyx b/python/cuml/cuml/manifold/umap/umap.pyx index f90673b697..41bdc51bd2 100644 --- a/python/cuml/cuml/manifold/umap/umap.pyx +++ b/python/cuml/cuml/manifold/umap/umap.pyx @@ -1307,6 +1307,7 @@ class UMAP(InteropMixin, CMajorInputTagMixin, SparseInputTagMixin, Base): knn_indices, knn_dists = extract_knn_graph( (knn_graph if knn_graph is not None else self.precomputed_knn), + X.shape[0], self._n_neighbors, mem_type=None, # mirrors the input graph mem type indices_dtype=("int32" if X_is_sparse else "int64"), diff --git a/python/cuml/cuml/manifold/utils.py b/python/cuml/cuml/manifold/utils.py index d33d8d9bbb..a49b8731b9 100644 --- a/python/cuml/cuml/manifold/utils.py +++ b/python/cuml/cuml/manifold/utils.py @@ -109,6 +109,7 @@ def _check_self_references(indices, orig_n_neighbors=None): def extract_knn_graph( knn_info, + n_samples, n_neighbors, mem_type="device", indices_dtype="int64", @@ -124,8 +125,10 @@ def extract_knn_graph( - Pairwise distances dense array of shape (n_samples, n_samples). - KNN graph sparse array. This is most efficient if the graph is in CSR format and contains 0 entries in `data` for all diagonal elements. + n_samples: int + Number of samples expected. n_neighbors: int - Number of nearest neighbors + Number of nearest neighbors required. mem_type : {"device", "host", None}, default="device" The desired output memory type. indices_dtype : dtype, default='int64' @@ -151,11 +154,11 @@ def extract_knn_graph( distances = check_array( distances, dtype="float32", order="C", mem_type=mem_type_init ) - if not indices.shape == distances.shape: + if indices.shape[0] != n_samples or indices.shape != distances.shape: raise ValueError( f"Expected indices and distances to have shape=(n_samples, " - f"n_neighbors), got indices.shape={indices.shape}, " - f"distances.shape={distances.shape}" + f"n_neighbors) where {n_samples=}, got " + f"indices.shape={indices.shape}, distances.shape={distances.shape}" ) if not _check_self_references(indices): raise ValueError( @@ -169,10 +172,13 @@ def extract_knn_graph( # Sparse KNN graph # - shape=(n_samples, n_samples) # - nnz=n_samples * orig_n_neighbors - if not (knn_info.ndim == 2 and knn_info.shape[0] == knn_info.shape[1]): + if not ( + knn_info.ndim == 2 + and knn_info.shape[0] == knn_info.shape[1] == n_samples + ): raise ValueError( - f"Expected a sparse array of shape=(n_samples, n_samples), " - f"got shape={knn_info.shape}" + f"Expected a sparse array of shape=(n_samples, n_samples) where " + f"{n_samples=}, got shape={knn_info.shape}" ) # Coerce to CSR. If the input was already CSR this is zero-copy and @@ -188,8 +194,6 @@ def extract_knn_graph( ) xp = cp if isinstance(knn_info.data, cp.ndarray) else np - n_samples = knn_info.shape[0] - orig_n_neighbors, remainder = divmod(knn_info.nnz, n_samples) if ( remainder == 0 @@ -245,12 +249,11 @@ def extract_knn_graph( knn_info = check_array( knn_info, dtype="float32", mem_type=mem_type_init ) - if knn_info.shape[0] != knn_info.shape[1]: + if not (knn_info.shape[0] == knn_info.shape[1] == n_samples): raise ValueError( - f"Expected a dense array of shape=(n_samples, n_samples), " - f"got shape={knn_info.shape}" + f"Expected a dense array of shape=(n_samples, n_samples) where " + f"{n_samples=}, got shape={knn_info.shape}" ) - n_samples = knn_info.shape[0] if n_samples < n_neighbors: raise ValueError( f"Precomputed KNN data requires n_samples >= n_neighbors. " diff --git a/python/cuml/tests/test_tsne.py b/python/cuml/tests/test_tsne.py index 9f247d7c1d..feae0c070d 100644 --- a/python/cuml/tests/test_tsne.py +++ b/python/cuml/tests/test_tsne.py @@ -61,41 +61,14 @@ def test_tsne_knn_graph_used( min_grad_norm=1e-12, ) - # Perform tsne with normal knn_graph + # Fit works and results in decent score with provided knn_graph Y = tsne.fit_transform(X, convert_dtype=True, knn_graph=knn_graph) + trust = trustworthiness(X, Y, n_neighbors=DEFAULT_N_NEIGHBORS) + assert trust >= 0.80 - trust_normal = trustworthiness(X, Y, n_neighbors=DEFAULT_N_NEIGHBORS) - - X_garbage = np.ones(X.shape) - knn_graph_garbage = neigh.kneighbors_graph( - X_garbage, mode="distance" - ).astype("float32") - - if type_knn_graph == "cuml": - knn_graph_garbage = cupyx.scipy.sparse.csr_matrix(knn_graph_garbage) - - tsne = TSNE( - random_state=1, - n_neighbors=DEFAULT_N_NEIGHBORS, - method=method, - perplexity=DEFAULT_PERPLEXITY, - learning_rate_method="none", - min_grad_norm=1e-12, - ) - - # Perform tsne with garbage knn_graph - Y = tsne.fit_transform(X, convert_dtype=True, knn_graph=knn_graph_garbage) - - trust_garbage = trustworthiness(X, Y, n_neighbors=DEFAULT_N_NEIGHBORS) - assert (trust_normal - trust_garbage) > 0.15 - - Y = tsne.fit_transform(X, convert_dtype=True, knn_graph=knn_graph_garbage) - trust_garbage = trustworthiness(X, Y, n_neighbors=DEFAULT_N_NEIGHBORS) - assert (trust_normal - trust_garbage) > 0.15 - - Y = tsne.fit_transform(X, convert_dtype=True, knn_graph=knn_graph_garbage) - trust_garbage = trustworthiness(X, Y, n_neighbors=DEFAULT_N_NEIGHBORS) - assert (trust_normal - trust_garbage) > 0.15 + # Fit errors if graph is bad + with pytest.raises(ValueError, match="Expected a sparse array of shape"): + tsne.fit_transform(X, knn_graph=knn_graph[:20, :20]) @pytest.mark.parametrize("type_knn_graph", ["cuml", "sklearn"]) diff --git a/python/cuml/tests/test_umap.py b/python/cuml/tests/test_umap.py index def7fb8968..39f55e3ece 100644 --- a/python/cuml/tests/test_umap.py +++ b/python/cuml/tests/test_umap.py @@ -630,7 +630,7 @@ def test_extract_knn_graph( # mem_type=None doesn't coerce mem_type indices, distances = extract_knn_graph( - knn_info, k, indices_dtype=indices_dtype, mem_type=None + knn_info, 30, k, indices_dtype=indices_dtype, mem_type=None ) assert indices.dtype == indices_dtype assert distances.dtype == "float32" @@ -641,7 +641,7 @@ def test_extract_knn_graph( np.testing.assert_array_equal(cp.asnumpy(indices), sol.indices) indices, distances = extract_knn_graph( - knn_info, k, indices_dtype=indices_dtype, mem_type="host" + knn_info, 30, k, indices_dtype=indices_dtype, mem_type="host" ) assert indices.dtype == indices_dtype assert distances.dtype == "float32" @@ -649,7 +649,7 @@ def test_extract_knn_graph( np.testing.assert_array_equal(indices, sol.indices) indices, distances = extract_knn_graph( - knn_info, k, indices_dtype=indices_dtype, mem_type="device" + knn_info, 30, k, indices_dtype=indices_dtype, mem_type="device" ) assert indices.dtype == indices_dtype assert distances.dtype == "float32" @@ -665,7 +665,12 @@ def test_extract_knn_graph_errors(): with pytest.raises( ValueError, match="Expected indices and distances to have shape" ): - extract_knn_graph(knn_info, 10) + extract_knn_graph(knn_info, 3, 2) + knn_info = (np.ones((2, 4)), np.ones((2, 4))) + with pytest.raises( + ValueError, match="Expected indices and distances to have shape" + ): + extract_knn_graph(knn_info, 3, 2) # tuple: n_neighbors > original_n_neighbors knn_info = tuple( @@ -682,7 +687,7 @@ def test_extract_knn_graph_errors(): "n_neighbors=25 was specified" ), ): - extract_knn_graph(knn_info, 25) + extract_knn_graph(knn_info, 30, 25) # tuple: expected self references knn_info = tuple( @@ -696,12 +701,15 @@ def test_extract_knn_graph_errors(): ValueError, match="Expected indices and distances to include self references", ): - extract_knn_graph(knn_info, 20) + extract_knn_graph(knn_info, 30, 20) # graph: invalid shape - knn_info = sp.random(10, 25, density=0.5) + knn_info = sp.random(29, 29, random_state=42, density=0.5) with pytest.raises(ValueError, match="Expected a sparse array of shape"): - extract_knn_graph(knn_info, 10) + extract_knn_graph(knn_info, 30, 10) + knn_info = sp.random(30, 29, random_state=42, density=0.5) + with pytest.raises(ValueError, match="Expected a sparse array of shape"): + extract_knn_graph(knn_info, 30, 10) # graph: invalid nnz knn_info = sp.csr_matrix( @@ -714,7 +722,7 @@ def test_extract_knn_graph_errors(): with pytest.raises( ValueError, match="Precomputed KNN graph has 4 nonzero elements" ): - extract_knn_graph(knn_info, 3) + extract_knn_graph(knn_info, 3, 3) # graph: n_neighbors > original_n_neighbors knn_info = kneighbors_graph(X, 20, include_self=True, mode="distance") @@ -725,12 +733,15 @@ def test_extract_knn_graph_errors(): "n_neighbors=25 was specified" ), ): - extract_knn_graph(knn_info, 25) + extract_knn_graph(knn_info, 30, 25) # pairwise: invalid shape knn_info = np.ones((5, 6)) with pytest.raises(ValueError, match="Expected a dense array of shape"): - extract_knn_graph(knn_info, 5) + extract_knn_graph(knn_info, 5, 5) + knn_info = np.ones((5, 5)) + with pytest.raises(ValueError, match="Expected a dense array of shape"): + extract_knn_graph(knn_info, 6, 5) # pairwise: n_neighbors > n_samples knn_info = np.ones((5, 5)) @@ -738,7 +749,7 @@ def test_extract_knn_graph_errors(): ValueError, match="Precomputed KNN data requires n_samples >= n_neighbors", ): - extract_knn_graph(knn_info, 10) + extract_knn_graph(knn_info, 5, 10) @pytest.mark.parametrize( @@ -762,7 +773,7 @@ def test_extract_knn_graph_zero_distance_rows(kind, mem_type): sol_indices = np.array([0, 3, 1, 1, 2, 3, 2, 1, 3, 3, 0, 1]) sol_distances = nn.kneighbors(X, 3)[0].flatten() - indices, distances = extract_knn_graph(graph, 3, mem_type=mem_type) + indices, distances = extract_knn_graph(graph, 4, 3, mem_type=mem_type) np.testing.assert_allclose(cp.asnumpy(distances), sol_distances, atol=1e-5) np.testing.assert_array_equal(cp.asnumpy(indices), sol_indices) @@ -792,7 +803,7 @@ def assert_same_memory(x, y): assert x_ptr == y_ptr indices, distances = extract_knn_graph( - graph, 20, mem_type=mem_type, indices_dtype=graph.indices.dtype + graph, 30, 20, mem_type=mem_type, indices_dtype=graph.indices.dtype ) assert_same_memory(indices, graph.indices) assert_same_memory(distances, graph.data)