refactor: expand CoreMLExportError to error on non-compatible configs for CoreML and remove more CoreML export tests (#42)

* refactor: centralize CoreML export compatibility validation

Weight/activation/LUT dtype checks and the activation-granularity check
were duplicated across the eager, graph, and palettization export paths.
Consolidate them into validate_coreml_compatibility() so CoreML's actual
export constraints live in one place, including the per-channel/per-block
activation restriction: an upstream coremltools MIL optimizer pass can
silently corrupt per-channel-quantized activations across reshape
boundaries, so CoreML export now rejects that configuration outright
instead of letting it through to a numerics-level failure.

* test: reject per-channel and per-block activation quantization on CoreML export

CoreML export now rejects these configs outright instead of letting them
through to a numerics-level SNR failure. Replace the xfail-based coverage
with direct CoreMLExportError assertions, and collapse the now-redundant
per-axis/per-dtype combinatorial sweeps since the rejection doesn't depend
on any of those variables.

* test: rename CoreML rejection assertion helper for dtype+granularity clarity

assert_coreml_finalize_rejects_unsupported_dtype and
COREML_DTYPE_REJECTION_MATCH are now used for granularity rejections too,
not just dtype ones, so their names no longer matched what they check.

Author: guru-desh

changelog.d/42.fixed

@@ -0,0 +1 @@
+Reject per-channel activation quantization on CoreML export

src/coreai_opt/_utils/export_utils.py

@@ -9,7 +9,9 @@
 import torch
 
 from coreai_opt._utils.torch_utils import is_tensor_on_cpu
-from coreai_opt.common import ExportBackend
+from coreai_opt.common import CoreMLExportError, ExportBackend
+from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization.spec.granularity import PerTensorGranularity, QuantizationGranularity
 
 COREML_SUPPORTED_WEIGHT_DTYPES: frozenset[torch.dtype] = frozenset(
     {
@@ -34,6 +36,47 @@
     }
 )
 
+COREML_SUPPORTED_ACTIVATION_GRANULARITIES: frozenset[type[QuantizationGranularity]] = frozenset(
+    {PerTensorGranularity}
+)
+
+_COREML_SUPPORTED_DTYPES_BY_TARGET: dict[CompressionTargetTensor, frozenset[torch.dtype]] = {
+    CompressionTargetTensor.WEIGHT: COREML_SUPPORTED_WEIGHT_DTYPES,
+    CompressionTargetTensor.ACTIVATION: COREML_SUPPORTED_ACTIVATION_DTYPES,
+    CompressionTargetTensor.LUT: COREML_SUPPORTED_LUT_DTYPES,
+}
+
+
+def validate_coreml_compatibility(
+    target: CompressionTargetTensor,
+    dtype: torch.dtype,
+    context: str,
+    granularity: QuantizationGranularity | None = None,
+) -> None:
+    """Raise CoreMLExportError if this weight/activation/LUT config isn't CoreML-exportable.
+
+    Centralizes every reason CoreML export can reject a quantization config, so
+    new restrictions are added here once rather than at each call site.
+
+    Args:
+        target (CompressionTargetTensor): Which tensor category is being checked.
+        dtype (torch.dtype): The quantization dtype to validate.
+        context (str): Human-readable description of what's being checked, used
+            in the error message (e.g. "weight 'conv.weight' of module 'conv'").
+        granularity (QuantizationGranularity | None): The quantization
+            granularity, if applicable. Only checked for ACTIVATION — CoreML
+            only supports per-tensor activation quantization.
+
+    Raises:
+        CoreMLExportError: If the dtype or granularity isn't supported.
+    """
+    if dtype not in _COREML_SUPPORTED_DTYPES_BY_TARGET[target]:
+        raise CoreMLExportError.from_dtype(dtype, context)
+    if target == CompressionTargetTensor.ACTIVATION and not isinstance(
+        granularity, tuple(COREML_SUPPORTED_ACTIVATION_GRANULARITIES)
+    ):
+        raise CoreMLExportError.from_config(granularity, context)
+
 
 def validate_mmap_backend_and_device(
     model: torch.nn.Module,

src/coreai_opt/common.py

@@ -163,8 +163,15 @@ class ExportBackend(_StrEnum, metaclass=_DeprecatedMemberEnumMeta):
 class CoreMLExportError(ValueError):
     """Raised when a model cannot be exported to the CoreML backend."""
 
-    def __init__(self, dtype: Any, context: str) -> None:
-        super().__init__(
-            f"CoreML export does not support dtype {dtype} on {context}. "
-            f"Use backend=ExportBackend.CoreAI instead."
-        )
+    def __init__(self, message: str) -> None:
+        super().__init__(f"{message} Use backend=ExportBackend.CoreAI instead.")
+
+    @classmethod
+    def from_dtype(cls, dtype: Any, context: str) -> CoreMLExportError:
+        """Build the error for an unsupported weight/activation/LUT dtype."""
+        return cls(f"CoreML export does not support dtype {dtype} on {context}.")
+
+    @classmethod
+    def from_config(cls, config: object, context: str) -> CoreMLExportError:
+        """Build the error for an unsupported quantization config attribute (e.g. granularity)."""
+        return cls(f"CoreML export does not support {type(config).__name__} on {context}.")

src/coreai_opt/palettization/kmeans/_prepare_for_export.py

@@ -12,16 +12,17 @@
 import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.export_utils import (
-    COREML_SUPPORTED_LUT_DTYPES,
     clear_parametrization_original as _clear_parametrization_original,
     prepare_mmap_dir as _prepare_mmap_dir,
+    validate_coreml_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 CoreMLExportError, ExportBackend
+from coreai_opt.common import ExportBackend
+from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.palettization.spec.fake_palettize import (
     _FakePalettizeImplBase,
 )
@@ -430,11 +431,11 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
         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:
-                if fake_palett_mod.lut_qspec.dtype not in COREML_SUPPORTED_LUT_DTYPES:
-                    raise CoreMLExportError(
-                        fake_palett_mod.lut_qspec.dtype,
-                        f"LUT of parameter '{param_name}' of module '{module_name}'",
-                    )
+                validate_coreml_compatibility(
+                    CompressionTargetTensor.LUT,
+                    fake_palett_mod.lut_qspec.dtype,
+                    f"LUT of parameter '{param_name}' of module '{module_name}'",
+                )
 
     _process_weight_palettization(model, backend=ExportBackend.CoreML)
 

src/coreai_opt/quantization/_eager/_prepare_for_mil_export.py

@@ -17,15 +17,11 @@
 import torch.nn.utils.parametrize as P
 from torch import nn
 
-from coreai_opt._utils.export_utils import (
-    COREML_SUPPORTED_ACTIVATION_DTYPES,
-    COREML_SUPPORTED_WEIGHT_DTYPES,
-)
+from coreai_opt._utils.export_utils import validate_coreml_compatibility
 from coreai_opt._utils.metadata_utils import CompressionType, MILCompressionMetadata
 from coreai_opt._utils.torch_utils import (
     get_parent_module_and_attr_name as _get_parent_module_and_attr_name,
 )
-from coreai_opt.common import CoreMLExportError
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._export_utils import (
     convert_dtype_for_torch_quantize as _convert_dtype_for_torch_quantize,
@@ -35,7 +31,6 @@
     validate_qformulation_for_mil_export as _validate_qformulation_for_mil_export,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
-from coreai_opt.quantization.spec.granularity import PerBlockGranularity
 
 
 def _process_weight_quantization(
@@ -96,22 +91,16 @@ def _process_activation_quantization(
 ) -> None:
     """Process activation quantization by replacing with Sequential quantize/dequantize.
 
-    Supports both per-tensor and per-channel quantization based on the
-    granularity axis of the fake quantization module.
+    By the time this runs, CoreML export compatibility has already been
+    validated, so only per-tensor activation granularity ever reaches here.
 
     Args:
         parent_module: The parent module containing the fake quantizer
         attr_name: The attribute name of the fake quantizer in the parent
         fake_quant_mod: The fake quantization module to replace
 
-    Raises:
-        ValueError: If the granularity is per-block (not supported for MIL
-            activation export).
-
     """
     _validate_qformulation_for_mil_export(fake_quant_mod)
-    if isinstance(fake_quant_mod.granularity, PerBlockGranularity):
-        raise ValueError("MIL export does not support per-block granularity for activations.")
 
     scale, zero_point, _ = _extract_quantization_params(fake_quant_mod)
     converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize(
@@ -151,24 +140,24 @@ def prepare_for_mil_export(model: nn.Module) -> nn.Module:
     """
     processed_fq_ids: set[int] = set()
 
-    # CoreML does not support certain dtypes (FP4, FP8, INT2, UINT2). Fail fast
-    # before mutating the model below.
+    # Fail fast if model is not coreml-exportable
     for module_name, module in model.named_modules():
         if P.is_parametrized(module):
             for param_name, parametrizations in module.parametrizations.items():
                 for p in parametrizations:
-                    if (
-                        _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT)
-                        and p.dtype not in COREML_SUPPORTED_WEIGHT_DTYPES
-                    ):
-                        raise CoreMLExportError(
-                            p.dtype, f"weight '{param_name}' of module '{module_name}'"
+                    if _is_module_fake_quant_target(p, CompressionTargetTensor.WEIGHT):
+                        validate_coreml_compatibility(
+                            CompressionTargetTensor.WEIGHT,
+                            p.dtype,
+                            f"weight '{param_name}' of module '{module_name}'",
                         )
-        if (
-            _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION)
-            and module.dtype not in COREML_SUPPORTED_ACTIVATION_DTYPES
-        ):
-            raise CoreMLExportError(module.dtype, f"activation quantizer of module '{module_name}'")
+        if _is_module_fake_quant_target(module, CompressionTargetTensor.ACTIVATION):
+            validate_coreml_compatibility(
+                CompressionTargetTensor.ACTIVATION,
+                module.dtype,
+                f"activation quantizer of module '{module_name}'",
+                module.granularity,
+            )
 
     for name, module in list(model.named_modules()):
         # Handle weight quantization parametrizations

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -17,18 +17,14 @@
 import torch
 from torch.fx import Node
 
-from coreai_opt._utils.export_utils import (
-    COREML_SUPPORTED_ACTIVATION_DTYPES,
-    COREML_SUPPORTED_WEIGHT_DTYPES,
-)
+from coreai_opt._utils.export_utils import validate_coreml_compatibility
 from coreai_opt._utils.fx_utils import get_node_type as _get_node_type
 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 (
     is_float4_dtype as _is_float4_dtype,
     sanitize_module_name as _sanitize_module_name,
 )
-from coreai_opt.common import CoreMLExportError
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization._export_utils import (
     canonicalize_qparam_shape as _canonicalize_qparam_shape,
@@ -45,7 +41,6 @@
     resolve_attr,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
-from coreai_opt.quantization.spec.granularity import PerBlockGranularity
 
 logger = logging.getLogger(__name__)
 
@@ -503,23 +498,17 @@ def _process_mil_activation_quantization(
 
     Converts activation FakeQuantize modules to Sequential modules containing
     quantize and dequantize operations that CoreMLTools can understand.
-    Supports both per-tensor (quantize_per_tensor) and per-channel
-    (quantize_per_channel) quantization based on the granularity axis.
-    Buffers are stored inside the modules.
+    Buffers are stored inside the modules. By the time this runs, CoreML
+    export compatibility has already been validated, so only per-tensor
+    activation granularity ever reaches here.
 
     Args:
         model: The graph module being modified
         fake_quant_node: The fake quantization node to process
         fake_quant_mod: The fake quantization module
 
-    Raises:
-        ValueError: If the granularity is per-block (not supported for MIL
-            activation export).
-
     """
     _validate_qformulation_for_mil_export(fake_quant_mod)
-    if isinstance(fake_quant_mod.granularity, PerBlockGranularity):
-        raise ValueError("MIL export does not support per-block granularity for activations.")
 
     scale, zero_point, _ = _extract_quantization_params(fake_quant_mod)
     converted_dtype, converted_zero_point = _convert_dtype_for_torch_quantize(
@@ -599,22 +588,22 @@ def prepare_for_mil_export(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         msg = "Model contains no fake quantization nodes"
         raise ValueError(msg)
 
-    # CoreML does not support certain dtypes. Fail fast before mutating the
-    # graph below.
+    # Fail fast if model is not coreml-exportable
     for fake_quant_node, fake_quant_mod in fake_quant_nodes:
         node_id = str(fake_quant_node.target)
         if _is_weight_fake_quant(fake_quant_node, fake_quant_mod):
-            if fake_quant_mod.dtype not in COREML_SUPPORTED_WEIGHT_DTYPES:
-                raise CoreMLExportError(
-                    fake_quant_mod.dtype,
-                    f"weight quantizer '{node_id}'",
-                )
+            validate_coreml_compatibility(
+                CompressionTargetTensor.WEIGHT,
+                fake_quant_mod.dtype,
+                f"weight quantizer '{node_id}'",
+            )
         else:
-            if fake_quant_mod.dtype not in COREML_SUPPORTED_ACTIVATION_DTYPES:
-                raise CoreMLExportError(
-                    fake_quant_mod.dtype,
-                    f"activation quantizer '{node_id}'",
-                )
+            validate_coreml_compatibility(
+                CompressionTargetTensor.ACTIVATION,
+                fake_quant_mod.dtype,
+                f"activation quantizer '{node_id}'",
+                fake_quant_mod.granularity,
+            )
 
     # Process all fake quantization nodes
     for fake_quant_node, fake_quant_mod in fake_quant_nodes:

tests/export/export_utils.py

@@ -29,9 +29,10 @@
 if platform.system() == "Darwin":
     from coreai.runtime import ComputeUnitKind, SpecializationOptions
 
-# Substring of the dtype guard message raised by the CoreML export validation. Shared so
-# test files asserting the rejection don't drift from one another.
-COREML_DTYPE_REJECTION_MATCH = "CoreML export does not support"
+# Substring of the guard message raised by CoreML export validation (dtype or
+# granularity/config rejection). Shared so test files asserting the rejection
+# don't drift from one another.
+COREML_REJECTION_MATCH = "CoreML export does not support"
 
 # Compute unit selection driven by the --compute-unit-kind pytest option (see
 # tests/conftest.py). Default is "interpreter" so a plain `pytest` run uses the
@@ -87,17 +88,18 @@ def _get_test_specialization_options() -> "SpecializationOptions | None":
     raise ValueError(msg)
 
 
-def assert_coreml_finalize_rejects_unsupported_dtype(finalizer: Any) -> None:
-    """Assert ``finalizer.finalize(backend=CoreML)`` rejects an unsupported dtype.
+def assert_coreml_finalize_rejects(finalizer: Any) -> None:
+    """Assert ``finalizer.finalize(backend=CoreML)`` rejects an unsupported config.
 
     CoreML does not support FP4, FP8, INT2, or UINT2 quantization or
-    palettization dtypes, so finalize must raise a ``CoreMLExportError`` rather
-    than emit an invalid model.
+    palettization dtypes, nor per-channel/per-block activation quantization
+    granularity, so finalize must raise a ``CoreMLExportError`` rather than
+    emit an invalid model.
 
     Args:
         finalizer (Any): A prepared ``Quantizer`` or ``KMeansPalettizer``.
     """
-    with pytest.raises(CoreMLExportError, match=COREML_DTYPE_REJECTION_MATCH):
+    with pytest.raises(CoreMLExportError, match=COREML_REJECTION_MATCH):
         finalizer.finalize(backend=ExportBackend.CoreML)
 
 

tests/export/test_eager_mil_export.py

@@ -94,6 +94,10 @@ def test_simple_model_export(
     parametrized_quant_config_general: ParametrizedQuantConfigs,
 ) -> None:
     """Test eager CoreML export with various quantization configurations."""
+    # CoreML rejects per-channel activation quantization outright.
+    if parametrized_quant_config_general.has_per_channel_activation_granularity:
+        pytest.skip("CoreML export rejects per-channel activation quantization")
+
     has_act_quant = parametrized_quant_config_general.has_activation_quantization
 
     _run_eager_mil_export_test(
@@ -114,6 +118,10 @@ def test_mnist_export(
     parametrized_quant_config_general: ParametrizedQuantConfigs,
 ) -> None:
     """Test eager CoreML export on MNIST model with various quantization configurations."""
+    # CoreML rejects per-channel activation quantization outright.
+    if parametrized_quant_config_general.has_per_channel_activation_granularity:
+        pytest.skip("CoreML export rejects per-channel activation quantization")
+
     has_act_quant = parametrized_quant_config_general.has_activation_quantization
 
     _run_eager_mil_export_test(
@@ -167,22 +175,31 @@ def test_gated_mlp_perchannel_act_export(
     parametrized_quant_config_perchannel_act_axis_coverage: ParametrizedQuantConfigs,
 ) -> None:
     """Test eager CoreML export with per-channel activation quantization axes.
-    Uses GatedMLPModel (uniform rank-3 activations throughout the model) to
-    test per-channel activation quantization across all valid axis values without
-    out-of-bounds errors.
+
+    Uses GatedMLPModel (uniform rank-3 activations throughout the model). CoreML
+    export rejects per-channel activation quantization outright; per-tensor cases
+    from this fixture should still export and verify normally.
     """
-    has_act_quant = (
-        parametrized_quant_config_perchannel_act_axis_coverage.has_activation_quantization
-    )
+    config = parametrized_quant_config_perchannel_act_axis_coverage
+
+    if config.has_per_channel_activation_granularity:
+        with pytest.raises(CoreMLExportError):
+            _run_eager_mil_export_test(
+                model=gated_mlp_model,
+                input_data=gated_mlp_model_input,
+                parametrized_quant_config=config,
+                expected_ops={},
+            )
+        return
 
     _run_eager_mil_export_test(
         model=gated_mlp_model,
         input_data=gated_mlp_model_input,
-        parametrized_quant_config=parametrized_quant_config_perchannel_act_axis_coverage,
+        parametrized_quant_config=config,
         expected_ops={
             "constexpr_blockwise_shift_scale": 3,
-            "quantize": 6 if has_act_quant else 0,
-            "dequantize": 6 if has_act_quant else 0,
+            "quantize": 6,
+            "dequantize": 6,
         },
     )