refactor: remove coremltools as required dependency and introduce numba kmeans1d#27
Open
guru-desh wants to merge 8 commits into
Open
refactor: remove coremltools as required dependency and introduce numba kmeans1d#27guru-desh wants to merge 8 commits into
numba kmeans1d#27guru-desh wants to merge 8 commits into
Conversation
guru-desh
commented
Jul 2, 2026
| raise ValueError("Cannot cluster an empty array.") | ||
| if not np.isfinite(values).all(): | ||
| raise ValueError("array must contain only finite values.") | ||
| k = min(k, n) |
Contributor
Author
There was a problem hiding this comment.
Currently, if k > n, the behavior is to get the minimum of the two. This is what the behavior was in coremltools's kmeans1d. We can choose to change it if wanted. I didn't change it because I wanted this PR to be as close to a pure refactor as possible
0ca13b4 to
c986b48
Compare
numba kmeans1d
4eafa01 to
e9631db
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR removes
coremltoolsas a required dependency forcoreai-opt.The only part of
coremltoolsthat was being used waskmeans1d. We have re-implementedkmeans1dincoreai-optusingnumbato keepcoreai-optplatform-agnostic.numbais a JIT compiler for Python that works by adding decorators to Python functions. This is different compared tocoremltoolskmeans1d that relied on AOT compilation of CPP code.The alternative of using
numbawould be to lazily do AOT compilation intorchusingtorch.utils.cpp_extension. This hasn't been tested.Performance Benchmarking
numbakmeans1dis around 1-3x faster thancoremltoolskmeans1d. This is mainly because coremltoolskmeans1ddoesn't compile with-O3. Once-O3is used, we should see the performance be the same according to thisnumbaforum post. AOT C++ is faster at small tensor sizes (10-20 elements) with lowkasnumbaJIT compilation time is much higher thankmeans1dexecution time.I have attached a markdown file where I ran performance comparisons of
numbaandkmeans1d:benchmark_results.md
Testing
kmeans1dis equivalent to coremltoolskmeans1d(tolerance is 1e-9)coremltools