Skip to content

refactor: remove coremltools as required dependency and introduce numba kmeans1d#27

Open
guru-desh wants to merge 8 commits into
apple:mainfrom
guru-desh:u/gurudesh/remove-coremltools-dependency
Open

refactor: remove coremltools as required dependency and introduce numba kmeans1d#27
guru-desh wants to merge 8 commits into
apple:mainfrom
guru-desh:u/gurudesh/remove-coremltools-dependency

Conversation

@guru-desh

@guru-desh guru-desh commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

Description

This PR removes coremltools as a required dependency for coreai-opt.

The only part of coremltools that was being used was kmeans1d. We have re-implemented kmeans1d in coreai-opt using numba to keep coreai-opt platform-agnostic.

numba is a JIT compiler for Python that works by adding decorators to Python functions. This is different compared to coremltools kmeans1d that relied on AOT compilation of CPP code.

The alternative of using numba would be to lazily do AOT compilation in torch using torch.utils.cpp_extension. This hasn't been tested.

Performance Benchmarking

numba kmeans1d is around 1-3x faster than coremltools kmeans1d. This is mainly because coremltools kmeans1d doesn't compile with -O3. Once -O3 is used, we should see the performance be the same according to this numba forum post. AOT C++ is faster at small tensor sizes (10-20 elements) with low k as numba JIT compilation time is much higher than kmeans1d execution time.

I have attached a markdown file where I ran performance comparisons of numba and kmeans1d:
benchmark_results.md

Testing

  1. We have equivalence tests to ensure that numba kmeans1d is equivalent to coremltools kmeans1d (tolerance is 1e-9)
  2. kmeans1d tests are copied over from coremltools
  3. These tests pass in CI

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@guru-desh guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency branch from 0ca13b4 to c986b48 Compare July 2, 2026 21:51
@guru-desh guru-desh changed the title build: remove coremltools as required dependency refactor: remove coremltools as required dependency and introduce numba kmeans1d Jul 2, 2026
@guru-desh guru-desh requested review from aseemw, pkmandke and u-simha July 2, 2026 21:54
@guru-desh guru-desh added the enhancement New feature or request label Jul 2, 2026
@guru-desh guru-desh marked this pull request as ready for review July 2, 2026 21:55
@guru-desh guru-desh force-pushed the u/gurudesh/remove-coremltools-dependency branch from 4eafa01 to e9631db Compare July 2, 2026 22:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants