Plumb metric and metric_kwds through to UMAP with nn_descent

[jcrist]

Previously we were erroneously missing this plumbing, leaving the metric and metric_arg as the defaults when nn_descent is used. In the long run we should redo how the parameters are passed when nn_descent is enabled to avoid this duplication (there's a few other uselessly exposed params like return_distances), but for now fixing the plumbing to be more correct seems fine.

On top of #6303, leaving as draft for now until that's merged.

Also re-enables tests on ARM, fixes #5441.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cjnolet]

/ok to test

[jcrist]

Rebased on main and marked ready for review.

[jcrist]

The actual fix is here (plumbing through the metric options to nn_descent_params). In the long run we should redesign the C++ layer to remove duplicate options - for now just ensuring they're forwarded correctly seems sufficient.

Everything else here is a simplification of the current pre-existing code.

[viclafargue]

LGTM, but we would have to make sure that NN Descent is actually compatible with all the metrics we make available. cc @jinsolp

[jcrist]

A quick browse through the code I believe all metrics are handled fine by nn_descent. Just to be sure I've also updated an existing metrics test to be parametrized on build_algo as well; this test runs fine locally as well.

The behavior before was incorrect (build_algo="nn_descent" would silently ignore the user selected metric), so I consider this PR strictly an improvement even in the off chance some parameter combo will error after this change. I'm fairly confident this won't happen though.

Planning to merge once CI is green.

[jinsolp]

Right now NN Descent in cuvs supports L2Expanded, CosineExpanded and InnerProduct, (link)
but in Raft Nn Descent supports only L2Expanded (link)

[jcrist]

/merge (just realized victor's review wasn't an actual approval)

[jcrist]

but in Raft Nn Descent supports only L2Expanded

Ah! Thanks @jinsolp! I'm curious - right now after plumbing the metric through nn_descent properly here I'm not seeing an error (and the added test passes), but you're right that it doesn't look like the metrics are supported after all. I don't see any validation in the raft code flow that would error on an invalid metric - perhaps the right additional fix here then is:

• Add validation on the python side for nn_descent to error on unsupported metrics
• Keep plumbing the args through properly, even though they appear to be mostly ignored on the raft side for now
• Keep the test, but update it to check for errors on unsupported metrics

Then whenever we get around to moving to cuvs instead we can increase the supported metrics pretty straightforwardly.

[jinsolp]

Yes I think that would be great! : )

[jcrist]

Looking at finishing this up today. @jinsolp, can you confirm that the existing support is for L2Expanded (called "sqeuclidean" in python) and not L2SqrtExpanded (called "euclidean" or "l2" in python)? Both raft and cuvs use else branches for this (and I had a hard time finding if any validation was done upstream to cull the set of DistanceTypes that may reach the branch).

If so, then validating would make build_algo="nn_descent" fail with the default configuration since we default to metric="euclidean". If this is an issue, would it be to be possible to get L2SqrtExpanded support into the raft (and cuvs) implementations this release to resolve this?

[jinsolp]

Right now with NN Descent raft, we have a distance epilogue that takes care of the sqrt part.
(So NN Descent in raft just calculates the L2Expanded, but it's sqrt-ed in the distance epilogue here)
The distance epilogue is made and passed on from the UMAP side (written here)

Summarizing; UMAP + NN Descent works with "euclidean" because UMAP's distance epilogue sqrt-s the L2Expanded in NN Descent (which is the only supported metric for NN Descent in raft as of now), making it a L2SqrtExpanded type.

Noticing that the distance epilogues are no longer in the cuvs version, we can add the L2SqrtExpanded type in cuvs to enable its use with cuML UMAP in the future.

[jcrist]

Hmmm, that seems a bit convoluted. For the cuvs implementation, I'd hope we could avoid the callback and just handle the different DistanceTypes natively (unless there's a good reason to do this with a user-provided callback instead).

As is, it looks like we always pass in the epilogue on the cuml umap side, which seems incorrect if we're computing L2Expanded instead of L2SqrtExpanded? I think this is wrong, happy to change it as part of this PR for now. As written, I guess this means we only support L2SqrtExpanded for nn_descent in cuml for now.

[jinsolp]

The default distance epilogue is set to return itself (i.e. is an identity op) declared here as part of the NN Descent types.
So it doesn't do anything (i.e. the distance is calculated as L2Expanded) if nothing specific is passed as the distance epilogue. UMAP inherits this epilogue type and overrides the () operator to return an sqrt-ed value of its input.

unless there's a good reason to do this with a user-provided callback instead

The distance epilogues are actually needed to make NN Descent work with HDBSCAN. We will need them in cuvs too eventually : ) (or find a smarter way to do this)

which seems incorrect if we're computing L2Expanded instead of L2SqrtExpanded?

You are right about this!

I guess this means we only support L2SqrtExpanded for nn_descent in cuml for now.

This should be fine for now 🙂

[jcrist]

Done - we should now:

• Properly forward the parameters through to libcuml/raft/etc...
• Error nicely on unsupported metrics for umap and friends (I consolidated the logic here to be more uniform too)

Should be ready for another round of review.

[jinsolp]

LGTM! Thank you for this PR!

[viclafargue]

LGTM! Just small change request

[csadorf]

One suggestion, but LGTM.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 67.39%. Comparing base (4384f17) to head (dfd6c65).
Report is 2 commits behind head on branch-25.04.

Additional details and impacted files

@@               Coverage Diff                @@
##           branch-25.04    #6304      +/-   ##
================================================
- Coverage         68.53%   67.39%   -1.14%     
================================================
  Files               204      204              
  Lines             13199    13230      +31     
================================================
- Hits               9046     8917     -129     
- Misses             4153     4313     +160     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jcrist]

/merge