Cleanup extract_knn_graph

[jcrist]

• 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. Part of Transition to new array input validation system #7428.
• 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.
• Adds support for specifying the output indices dtype, reducing memory usage and copies for sparse data code paths.
• 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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e3a26000-e76e-4fd0-aad5-8668509a322a

📥 Commits

Reviewing files that changed from the base of the PR and between 8cc749d and c7fbee9.

📒 Files selected for processing (5)

• python/cuml/cuml/manifold/t_sne.pyx
• python/cuml/cuml/manifold/umap/umap.pyx
• python/cuml/cuml/manifold/utils.py
• python/cuml/tests/test_tsne.py
• python/cuml/tests/test_umap.py

🚧 Files skipped from review as they are similar to previous changes (4)

• python/cuml/cuml/manifold/t_sne.pyx
• python/cuml/cuml/manifold/umap/umap.pyx
• python/cuml/tests/test_umap.py
• python/cuml/cuml/manifold/utils.py

---

📝 Walkthrough

Summary by CodeRabbit

Release Notes

• Refactor

  • Consolidated KNN graph extraction/normalization into a single shared path for t-SNE and UMAP.
  • Improved handling of precomputed KNN inputs (tuple, sparse CSR/CSC/COO, and dense pairwise distances), including more robust dtype and memory-type behavior.

• Documentation

  • Updated t-SNE and UMAP documentation to clarify accepted precomputed KNN formats and that knn_graph takes precedence.

• Tests

  • Expanded and strengthened UMAP and t-SNE test coverage for KNN extraction correctness, dtype placement, edge cases, and validation errors.

Walkthrough

Relocates KNN graph extraction utilities from cuml.common.sparsefuncs to cuml.manifold.utils with CUDA-accelerated validation kernels and flexible dtype/memory handling. Updates TSNE and UMAP callers to import and invoke the new API with conditional dtype and memory-type routing. Adds comprehensive parametrized test coverage including error handling, zero-copy validation, and sparse import normalization.

Changes

KNN Graph Extraction Relocation

Layer / File(s)	Summary
New extract_knn_graph implementation in manifold/utils python/cuml/cuml/manifold/utils.py	Introduces CUDA validation kernels _is_axis_1_sorted_kernel, _check_indices_per_row, _check_distances_sorted, and _check_self_references to verify CSR structure and neighbor order. Implements extract_knn_graph supporting tuple (indices, distances), sparse CSR KNN graphs, and dense pairwise distance matrices, with dtype coercion, neighbor extraction via argpartition, sorting, trimming, and output memory-type normalization.
Remove old utilities from sparsefuncs python/cuml/cuml/common/sparsefuncs.py	Removes extract_knn_graph, extract_sparse_knn_graph, extract_pairwise_dists, and _determine_k_from_arrays; updates copyright year and drops SciPy sparse and Cuml input-conversion imports.
TSNE caller update python/cuml/cuml/manifold/t_sne.pyx	Switches extract_knn_graph import from cuml.common.sparsefuncs to cuml.manifold.utils. Updates docstrings for precomputed_knn and knn_graph parameters to enumerate accepted forms. Refactors TSNE.fit to invoke extract_knn_graph with indices_dtype conditional on sparse input, using returned arrays directly for device-pointer extraction instead of prior explicit CuPy conversion and int32 casting.
UMAP caller update python/cuml/cuml/manifold/umap/umap.pyx	Relocates extract_knn_graph import to cuml.manifold.utils with reorganized cimports. Updates precomputed_knn, knn_graph, and fit_transform docstrings to describe supported input forms. Modifies _attrs_to_cpu to use cp.asnumpy(...) instead of to_cpu(...) helper. Refactors UMAP.fit to pass mem_type=None with conditional indices_dtype, and branches pointer derivation for CuPy device vs host arrays, removing prior forced conversions.
Tests and import cleanup python/cuml/tests/test_umap.py, python/cuml/tests/test_tsne.py	Normalizes sparse imports to cp_sp/sp, enabling uniform backend selection by input_type. Updates existing sparse-matrix construction in test_umap_transform_on_digits_sparse. Adds parametrized test_extract_knn_graph for memory types/dtypes/input formats, test_extract_knn_graph_errors for validation failures, test_extract_knn_graph_zero_distance_rows for expected outputs, and test_extract_knn_graph_from_kneighbors_graph_zero_copy for no-copy verification. Updates sparse references and error message assertions. Simplifies test_tsne_knn_graph_used to validate trustworthiness and assert shape-mismatch errors.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• rapidsai/cuml#7915: Aligns with UMAP precomputed knn_graph handling refactoring to support host-provided KNN data via memory-type parameter routing instead of forced CuPy conversions.

Suggested labels

improvement

Suggested reviewers

• dantegd
• lowener
• csadorf

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 53.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Cleanup extract_knn_graph' accurately and concisely describes the PR's primary action: refactoring and improving the extract_knn_graph function implementation.
Description check	✅ Passed	The description comprehensively explains the changes and aligns with the file modifications: moving the function to manifold utils, fixing bugs, adding features, and improving validation and documentation.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/cuml/cuml/manifold/utils.py`:
- Around line 178-188: The current validation in the error check only rejects
cases where n_neighbors is strictly greater than n_samples, but it should also
reject the edge case where n_neighbors equals n_samples. When n_neighbors ==
n_samples, the argpartition operation will select all columns including the
self-index (diagonal element) even after filling it with infinity, resulting in
self-neighbors in the KNN graph. Modify the condition in the ValueError check to
use n_samples <= n_neighbors instead of n_samples < n_neighbors to properly
reject this edge case and prevent invalid self-neighbors in the output KNN
graph.
- Around line 115-128: The current validation only checks if total non-zero
elements divide evenly by n_samples, but doesn't verify that each row has the
same number of neighbors. This allows uneven rows to pass validation and get
incorrectly reshaped, mixing neighbor entries across samples. After converting
to CSR format with tocsr() and before reshaping knn_info.indices, add validation
to confirm every row has exactly orig_n_neighbors elements by checking the
differences between consecutive indptr values in the CSR matrix are all equal to
orig_n_neighbors, and raise a ValueError if any row has a different count.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 38dbfa58-f384-455c-915f-caf77b3feaa7

📥 Commits

Reviewing files that changed from the base of the PR and between bac16da and 8d99652.

📒 Files selected for processing (5)

• python/cuml/cuml/common/sparsefuncs.py
• python/cuml/cuml/manifold/t_sne.pyx
• python/cuml/cuml/manifold/umap/umap.pyx
• python/cuml/cuml/manifold/utils.py
• python/cuml/tests/test_umap.py

[jcrist]

Hmmm, it's not clear to me whether the precomputed KNN data should include self references (i.e. should data on a sample indicate the closest sample is itself)? Upstream UMAP wants that:

The nearest neighbor of each item should be itself, e.g. the nearest neighbor of item 0 should be 0, the nearest neighbor of item 1 is 1 and so on. Please note that you will not be able to transform new data in this case.

What's the expected behavior for our implementation? Do we require self references? What happens if they're not there (if we need them) / if they're there (if we don't need them)?

Getting self references from a KNN graph input is a bit tricky, since the expected values are 0 and would thus be dropped by a canonical sparse array (they're also not included by default in kneighbors_graph, you have to also pass in include_self=True). Requiring them makes it less likely we can use input data from a sparse KNN graph without requiring a copy.

FWIW - all tests in test_umap.py before this PR did provide data with self references, so my assumption is they're desired. Just want to clarify whether that's true or not before I fix things up. cc @jinsolp for advice.

[jinsolp]

Yes current implementation requires self-references in the knn graph.
Right now if they're not there, we would silently skip the first NN information.

However, it's not an algorithmic requirement, so we can update a few pieces in the current UMAP code to remove this requirement if needed(e.g. the smooth_knn_dist_kernel psum loop) :

https://github.com/rapidsai/cuml/blob/427c2cf621bef07cdc1e842eb5b4c7ecec28c021/cpp/src/umap/fuzzy_simpl_set/naive.cuh#L140

[jcrist]

Thanks Jinsol!

I think we can leave the C++ layer alone - requiring self references is fine (it's what upstream umap does).

In this PR I'll:

• Update the code to always produce self references
• Include a cheap check for this to ensure we don't get invalid inputs. We can check this easily without requiring any memory transfers, so I'm not worried about added costs.
• Update the docstrings to note the requirement.

[dantegd]

Just have two questions, besides that it looks good to me

[jcrist]

/merge