refactor: remove coremltools dependency and use torch.utils.cpp_extension to load kmeans1d

.gitignore

@@ -47,7 +47,6 @@ uv.lock
 *.nc
 .python-version
 /.tox
-deps/
 scratch/
 /local_stash
 /model_weights

.pre-commit-config.yaml

@@ -23,6 +23,7 @@ repos:
         language: system
         entry: scripts/pre_commit/add_license_header.py --license-file configs/BSD-3-LICENSE-HEADER-TEMPLATE --start-year 2026
         files: '\.(py|sh|js|css|html)$|(^|/)(GNUmakefile|[Mm]akefile)$'
+        exclude: ^src/coreai_opt/deps/
 
   # ----------------------------------------------------------------------------
   # 0.2 Update LICENSE year

NOTICE.txt

@@ -0,0 +1,25 @@
+Copyright 2026 Apple Inc.
+
+This project contains content adapted from kmeans1d (https://github.com/dstein64/kmeans1d), the license for which follows:
+
+MIT License
+
+Copyright (c) 2019 Daniel Steinberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

changelog.d/31.changed

@@ -0,0 +1 @@
+Replace the coremltools-based 1D k-means used by palettization with a vendored C++ core that is JIT-compiled at runtime via `torch.utils.cpp_extension`. `coremltools` is no longer a runtime dependency (it is now an optional dependency, installable via the `coreml` extra). This requires a C++ compiler to be available on the host at runtime.

pyproject.toml

@@ -31,11 +31,16 @@ classifiers = [
 ]
 dynamic = [ "version" ]
 dependencies = [
-    "coremltools>=8.3",
-    "numpy>=2,<2.4",              # TODO: Remove once coremltools has released a version > 9.0
+    # Required at runtime by torch.utils.cpp_extension to JIT-compile the vendored
+    # kmeans1d C++ core (a C++ toolchain must also be present on the host).
+    "ninja>=1.11",
+    "numpy>=2",
     "pydantic>=2.0.0",
+    "pyyaml>=6.0",
     "rich>=13.0.0",
     "safetensors>=0.5.3,<=0.7.0",
+    # Required at runtime by torch.utils.cpp_extension
+    "setuptools>=42",
     # PyTorch >= 2.9.0 requires torchao >= 0.15.0
     # Python's standard dependency specification (PEP 508) doesn't support conditional dependencies
     # based on other package versions. We can either 1) add a stricter check to require the newer torchao
@@ -54,6 +59,10 @@ coreai = [
     "coreai-torch==0.4.1",
     "scikit-learn>=1.7.2",
 ]
+coreml = [
+    "coremltools>=8.3",
+    "numpy>=2,<2.4",    # TODO: Remove once coremltools has released a version > 9.0
+]
 [project.urls]
 Changelog = "https://github.com/apple/coreai-optimization/blob/main/CHANGELOG.md"
 Repository = "https://github.com/apple/coreai-optimization"
@@ -95,7 +104,9 @@ docs = [
 # installs CoreAI by default (see tool.uv.default-groups). The extra is the
 # single source of truth; this group self-references it to avoid duplicating the
 # package list.
+# This applies to the `coreml` group as well.
 coreai = [ "coreai-opt[coreai]" ]
+coreml = [ "coreai-opt[coreml]" ]
 # Used in CI to force latest mimimum supported torch version
 # These torch versions must be in bounds of torch versions listed in project dependencies
 highest_tested_torch = [
@@ -145,6 +156,7 @@ tutorial = [
     "papermill>=2.7.0",
     "torchinfo>=1.8.0",
     { include-group = "coreai" },
+    { include-group = "coreml" },
     { include-group = "torchvision" },
 ]
 
@@ -153,14 +165,14 @@ package-dir = { "" = "src" }
 [tool.setuptools.dynamic]
 version = { attr = "coreai_opt._about.__version__" }
 [tool.setuptools.package-data]
-coreai_opt = [ "py.typed" ]
+coreai_opt = [ "deps/_kmeans1d/_core.cpp", "deps/_kmeans1d/LICENSE", "py.typed" ]
 [tool.setuptools.packages]
 find.exclude = [ "coreai_opt_benchmarking*", "tests*" ]
 find.where = [ "src" ]
 
 # make env group installed by default with uv sync
 [tool.uv]
-default-groups = [ "dev", "coreai" ]
+default-groups = [ "dev", "coreai", "coreml" ]
 index = [
     { explicit = true, name = "pytorch-cpu", url = "https://download.pytorch.org/whl/cpu" },
     { explicit = true, name = "pytorch-cu128", url = "https://download.pytorch.org/whl/cu128" },

src/coreai_opt/coreai_utils/_utils/palettize_utils.py

@@ -24,19 +24,12 @@
 )
 from coreai_opt.coreai_utils._utils.graph_utils import _select_input_output_channel_axis
 from coreai_opt.coreai_utils.common import CompressionGranularity as _CompressionGranularity
+from coreai_opt.deps import _kmeans1d
 
 logger = logging.getLogger(__name__)
 
 _CONV2D_OP = "coreai.conv2d"
 
-try:
-    import kmeans1d as _kmeans1d
-
-    _HAS_KMEANS1D = True
-except ImportError:
-    _kmeans1d = None  # type: ignore[assignment]
-    _HAS_KMEANS1D = False
-
 LutParams = namedtuple("LutParams", "indices lut vector_axis")
 
 _SUPPORTED_NBITS: tuple[int, ...] = (1, 2, 3, 4, 6, 8)
@@ -99,8 +92,7 @@ def _get_kmeans_lookup_table_and_weight(
         weight.shape[1] == 1 and num_weights >= 10_000 and weight.dtype == np.float16
     )
 
-    if (is_better_to_use_kmeans1d and _HAS_KMEANS1D) or force_kmeans1d:
-        assert _HAS_KMEANS1D, "Unable to import kmeans1d, please make sure it's installed."
+    if is_better_to_use_kmeans1d or force_kmeans1d:
         values, indices, counts = np.unique(weight, return_inverse=True, return_counts=True)
         indices = indices.flatten()
         n_clusters = min(len(values), lut_len)

src/coreai_opt/deps/_kmeans1d/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Daniel Steinberg
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

src/coreai_opt/deps/_kmeans1d/__init__.py

@@ -0,0 +1,27 @@
+# MIT License
+#
+# Copyright (c) 2019 Daniel Steinberg
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+# Copyright © 2026 Apple Inc.
+
+from coreai_opt.deps._kmeans1d.core import Clustered, cluster
+
+__all__ = ["Clustered", "cluster"]