Expose kmeans to python #729
Conversation
benfred
added
improvement
|
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. |
|
@benfred this looks great, but one of the things we're being asked for quite a bit today is to expose the hierarchical kmeans to Python. Any chance we can also expose those functions? I don't mind doing it as a follow-up, given that this PR is already feature complete. |
|
/ok to test |
cpp/src/cluster/kmeans.cuh
Outdated
|
|
||
| rmm::device_uvector<char> workspace(n_samples * sizeof(IndexT), stream); | ||
|
|
||
| rmm::device_uvector<DataT> x_norms(n_samples, stream); |
There was a problem hiding this comment.
Can the newer mdarray/mdspan API be used here? For the allocation of memory and the calls to raft functions that accept it.
There was a problem hiding this comment.
I've used the newer mdarray functions where possible in the last commit (but there are some cases where a device_uvector is expected, like the workspace etc, so I've left those as is)
No description provided.