Skip to content

Cleanup extract_knn_graph #8272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 3 commits into from
Jun 23, 2026
Merged

Cleanup extract_knn_graph #8272

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c975a90
Cleanup `extract_knn_graph`
jcrist Jun 16, 2026
8cc749d
Respond to feedback
jcrist Jun 23, 2026
c7fbee9
Add check for correct `n_samples`, and fix TSNE test
jcrist Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 1 addition & 214 deletions python/cuml/cuml/common/sparsefuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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
68 changes: 36 additions & 32 deletions python/cuml/cuml/manifold/t_sne.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -610,20 +620,14 @@ 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 = <uintptr_t>knn_dists_cp.data.ptr
knn_indices_ptr = <uintptr_t>knn_indices_cp.data.ptr
knn_indices, knn_dists = extract_knn_graph(
knn_graph,
n_samples,
params.n_neighbors,
indices_dtype="int32" if sparse_fit else "int64",
)
knn_dists_ptr = <uintptr_t>knn_dists.data.ptr
knn_indices_ptr = <uintptr_t>knn_indices.data.ptr

# Allocate output array
embedding = cupy.zeros(
Expand Down
Loading
Loading