Add pairwise_distance api's for C, Python and Rust#142
Add pairwise_distance api's for C, Python and Rust#142rapids-bot[bot] merged 6 commits intorapidsai:branch-24.06from
Conversation
benfred
added
improvement
cjnolet
left a comment
cjnolet
left a comment
There was a problem hiding this comment.
I think overall this PR looks great. I'm pre-approving, but do want to make sure we verify the behavior of the C api for unsupported type combinations. Maybe you've already done this, but I just want to avoid any unintended behaviors.
| * @param[in] metric distance to evaluate | ||
| * @param[in] metric_arg metric argument (used for Minkowski distance) | ||
| */ | ||
| cuvsError_t cuvsPairwiseDistance(cuvsResources_t res, |
There was a problem hiding this comment.
This looks so simple and elegant to use!
| auto y_mds = cuvs::core::from_dlpack<mdspan_type>(y_tensor); | ||
| auto distances_mds = cuvs::core::from_dlpack<distances_mdspan_type>(distances_tensor); | ||
|
|
||
| cuvs::distance::pairwise_distance(*res_ptr, x_mds, y_mds, distances_mds, metric, metric_arg); |
There was a problem hiding this comment.
Do we know that is always going to use the float (instantiated) version?
| distances_tensor.dl_tensor.strides = NULL; | ||
|
|
||
| // run pairwise distances | ||
| cuvsPairwiseDistance(res, &dataset_tensor, &queries_tensor, &distances_tensor, metric, 2.0); |
There was a problem hiding this comment.
If you haven't done it before, can you try this w/ double? (Just making sure we don't get any unnexpected surprises from the UX or the compile times / binary size).
| from pylibraft.common.cai_wrapper import wrap_array | ||
| from pylibraft.common.interruptible import cuda_interruptible | ||
| from pylibraft.distance.pairwise_distance import DISTANCE_TYPES | ||
| from pylibraft.neighbors.common import _check_input_array |
There was a problem hiding this comment.
We'll probably need to bring this over at some point too (no rush for 24.06 since we already depend on pylibraft.
|
/merge |
2 participants
No description provided.