fix: raise CoreMLExportError on incompatible palettization configs

src/coreai_opt/_utils/export_utils.py

@@ -11,6 +11,7 @@
 from coreai_opt._utils.torch_utils import is_tensor_on_cpu
 from coreai_opt.common import CoreMLExportError, ExportBackend
 from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization.spec import QuantizationSpec
 from coreai_opt.quantization.spec.granularity import PerTensorGranularity, QuantizationGranularity
 
 COREML_SUPPORTED_WEIGHT_DTYPES: frozenset[torch.dtype] = frozenset(
@@ -78,6 +79,52 @@ def validate_coreml_compatibility(
         raise CoreMLExportError.from_config(granularity, context)
 
 
+def validate_coreml_palettization_compatibility(
+    cluster_dim: int,
+    lut_qspec: QuantizationSpec | None,
+    enable_per_channel_scale: bool,
+    context: str,
+) -> None:
+    """Raise CoreMLExportError if this palettization config isn't CoreML-exportable.
+
+    Checks the LUT dtype (delegating to validate_coreml_compatibility) and
+    whether the config combines multiple features CoreML/MIL export cannot yet
+    fuse into a single compatible op chain: CoreML export supports at most one
+    of {vector palettization, LUT quantization, per-channel scale} at a time;
+    combining any two hits an unsupported CoreML/MIL op configuration
+    (mismatched tensor ranks, or `lut_to_dense` divisibility errors).
+
+    Args:
+        cluster_dim (int): Palettization cluster dimension; > 1 indicates
+            vector palettization.
+        lut_qspec (QuantizationSpec | None): LUT quantization spec, or None if
+            the LUT is not quantized.
+        enable_per_channel_scale (bool): Whether per-channel scaling is enabled.
+        context (str): Human-readable description of what's being checked.
+
+    Raises:
+        CoreMLExportError: If the LUT dtype isn't supported, or if two or more
+            of the three features above are combined.
+    """
+    if lut_qspec is not None:
+        validate_coreml_compatibility(
+            CompressionTargetTensor.LUT, lut_qspec.dtype, f"LUT of {context}"
+        )
+
+    active_features = []
+    if cluster_dim > 1:
+        active_features.append("cluster_dim")
+    if lut_qspec is not None:
+        active_features.append("lut_qspec")
+    if enable_per_channel_scale:
+        active_features.append("enable_per_channel_scale")
+
+    if len(active_features) >= 2:
+        raise CoreMLExportError(
+            f"CoreML export does not support {' + '.join(active_features)} on {context}."
+        )
+
+
 def validate_mmap_backend_and_device(
     model: torch.nn.Module,
     backend: ExportBackend,

src/coreai_opt/palettization/kmeans/_prepare_for_export.py

@@ -14,15 +14,14 @@
 from coreai_opt._utils.export_utils import (
     clear_parametrization_original as _clear_parametrization_original,
     prepare_mmap_dir as _prepare_mmap_dir,
-    validate_coreml_compatibility,
+    validate_coreml_palettization_compatibility,
 )
 from coreai_opt._utils.import_utils import lazy_import_coreai_torch
 from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata
 from coreai_opt._utils.torch_utils import (
     mmap_module_state_dict as _mmap_module_state_dict,
 )
 from coreai_opt.common import ExportBackend
-from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.palettization.spec.fake_palettize import (
     _FakePalettizeImplBase,
 )
@@ -430,12 +429,16 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
             continue
         for param_name, parametrizations in module.parametrizations.items():
             _, fake_palett_mod = _find_fake_palett_parametrization(parametrizations)
-            if fake_palett_mod is not None and fake_palett_mod.lut_qspec is not None:
-                validate_coreml_compatibility(
-                    CompressionTargetTensor.LUT,
-                    fake_palett_mod.lut_qspec.dtype,
-                    f"LUT of parameter '{param_name}' of module '{module_name}'",
-                )
+            if fake_palett_mod is None:
+                continue
+
+            context = f"parameter '{param_name}' of module '{module_name}'"
+            validate_coreml_palettization_compatibility(
+                fake_palett_mod.cluster_dim,
+                fake_palett_mod.lut_qspec,
+                fake_palett_mod.enable_per_channel_scale,
+                context,
+            )
 
     _process_weight_palettization(model, backend=ExportBackend.CoreML)
 

tests/export/test_kmeans_export.py

@@ -63,15 +63,17 @@ def _has_float_lut(config: ParametrizedPalettConfigs) -> bool:
     return config.lut_qspec is not None and config.lut_qspec.dtype.is_floating_point
 
 
-def _assert_coreml_rejects_unsupported_lut(
+def _assert_coreml_rejects(
     model: torch.nn.Module,
     input_data: torch.Tensor,
     config: KMeansPalettizerConfig,
 ) -> None:
-    """Assert finalize(CoreML) rejects an unsupported LUT dtype.
+    """Assert finalize(CoreML) rejects an unsupported palettization config.
 
-    CoreML/MIL does not support FP or INT2 LUT quantization, so finalize must raise
-    rather than emit an invalid model.
+    CoreML/MIL rejects certain LUT dtypes and certain combinations of
+    palettization features (vector palettization, LUT quantization,
+    per-channel scale), so finalize must raise rather than emit an invalid
+    model.
     """
     model.eval()
     palettizer = KMeansPalettizer(model, config)
@@ -95,32 +97,17 @@ def _skip_heavy_mnist_configs(config: ParametrizedPalettConfigs) -> None:
             pytest.skip(f"MNIST only tests lut_qspec with int8 and float8_e4m3fn, got {dtype}")
 
 
-def _skip_unsupported_mil_configs(
-    backend: ExportBackend,
-    config: ParametrizedPalettConfigs,
-) -> None:
-    """Skip CoreML configs with unsupported feature combinations."""
-    if backend != ExportBackend.CoreML:
-        return
-
+def _has_unsupported_mil_combo(config: ParametrizedPalettConfigs) -> bool:
+    """Whether the config combines >=2 of {vector palettization, LUT
+    quantization, per-channel scale} -- verified to be the exact set
+    CoreML/MIL cannot export: any single one of these works in isolation, but
+    combining two or more fails with a rank-mismatch, missing-vector_axis, or
+    LUT-divisibility error depending on the pair.
+    """
     is_vector = config.cluster_dim > 1
     has_lut_quant = config.lut_qspec is not None
     has_pcs = config.enable_per_channel_scale
-
-    # Vector palettization + LUT quantization
-    if is_vector and has_lut_quant:
-        # TODO: add CoreML export support for palettization combos.
-        pytest.skip("CoreML export not supported for vector palettization + LUT quantization.")
-
-    # Vector palettization + per-channel scale
-    if is_vector and has_pcs:
-        # TODO: add CoreML export support for palettization combos.
-        pytest.skip("CoreML export not supported for vector palettization + per-channel scale.")
-
-    # LUT quantization + per-channel scale
-    if has_lut_quant and has_pcs:
-        # TODO: add CoreML export support for palettization combos.
-        pytest.skip("CoreML export not supported for LUT quantization + per-channel scale.")
+    return sum([is_vector, has_lut_quant, has_pcs]) >= 2
 
 
 @pytest.mark.parametrize("backend", [ExportBackend.CoreML, ExportBackend.CoreAI])
@@ -134,12 +121,13 @@ def test_simple_model_export(
     config = parametrized_palett_config.config
     granularity = parametrized_palett_config.granularity
 
-    if backend == ExportBackend.CoreML and _has_float_lut(parametrized_palett_config):
-        _assert_coreml_rejects_unsupported_lut(simple_conv_linear_model, simple_model_input, config)
+    if backend == ExportBackend.CoreML and (
+        _has_float_lut(parametrized_palett_config)
+        or _has_unsupported_mil_combo(parametrized_palett_config)
+    ):
+        _assert_coreml_rejects(simple_conv_linear_model, simple_model_input, config)
         return
 
-    _skip_unsupported_mil_configs(backend, parametrized_palett_config)
-
     if (
         backend == ExportBackend.CoreML
         and parametrized_palett_config.cluster_dim > 1
@@ -175,12 +163,13 @@ def test_mnist_export(
     config = parametrized_palett_config.config
     granularity = parametrized_palett_config.granularity
 
-    if backend == ExportBackend.CoreML and _has_float_lut(parametrized_palett_config):
-        _assert_coreml_rejects_unsupported_lut(custom_test_mnist_model, mnist_example_input, config)
+    if backend == ExportBackend.CoreML and (
+        _has_float_lut(parametrized_palett_config)
+        or _has_unsupported_mil_combo(parametrized_palett_config)
+    ):
+        _assert_coreml_rejects(custom_test_mnist_model, mnist_example_input, config)
         return
 
-    _skip_unsupported_mil_configs(backend, parametrized_palett_config)
-
     # The MNIST model has 6 weight-bearing layers (conv1, conv2, conv_transpose1,
     # conv_transpose2, dense1, dense2). For axis=1 with group_size=2, conv1's
     # axis-1 (in_channels=1) is not divisible, so palettization is skipped there.