refactor: add InvalidExecutionModeError for unknown execution modes

src/coreai_opt/inspection/model_inspector.py

@@ -18,7 +18,10 @@
 from coreai_opt.base_model_compressor import _BaseModelCompressor
 from coreai_opt.palettization import KMeansPalettizer
 from coreai_opt.quantization import Quantizer
-from coreai_opt.quantization.config.quantization_config import ExecutionMode
+from coreai_opt.quantization.config.quantization_config import (
+    ExecutionMode,
+    InvalidExecutionModeError,
+)
 
 from ._eager_mode import parse_ops_for_eager as _parse_ops_for_eager
 from ._formatting import format_model_summary as _format_model_summary
@@ -138,8 +141,7 @@ def _validate_args(
             raise TypeError(msg)
 
         if execution_mode not in (ExecutionMode.GRAPH, ExecutionMode.EAGER):
-            msg = f"Unknown execution_mode {execution_mode}. Expected 'graph' or 'eager'."
-            raise ValueError(msg)
+            raise InvalidExecutionModeError(execution_mode)
 
         if example_inputs is None and not (
             isinstance(model, torch.fx.GraphModule) and execution_mode == ExecutionMode.GRAPH

src/coreai_opt/quantization/__init__.py

@@ -15,13 +15,15 @@
 # Import config after spec to avoid circular imports
 from .config import (
     ExecutionMode,
+    InvalidExecutionModeError,
     ModuleQuantizerConfig,
     QuantizerConfig,
 )
 from .quantizer import Quantizer
 
 __all__ = [
     "ExecutionMode",
+    "InvalidExecutionModeError",
     "ModuleQuantizerConfig",
     "QuantizationSpec",
     "Quantizer",

src/coreai_opt/quantization/config/__init__.py

@@ -7,6 +7,7 @@
 
 from .quantization_config import (
     ExecutionMode,
+    InvalidExecutionModeError,
     KVCacheQuantConfig,
     ModuleQuantizerConfig,
     OpQuantizerConfig,
@@ -16,6 +17,7 @@
 
 __all__ = [
     "ExecutionMode",
+    "InvalidExecutionModeError",
     # Configuration classes
     "KVCacheQuantConfig",
     "ModuleQuantizerConfig",

src/coreai_opt/quantization/config/quantization_config.py

@@ -158,6 +158,13 @@ class ExecutionMode(_StrEnum, metaclass=_DeprecatedMemberEnumMeta):
         """Deprecated. Use ``ExecutionMode.GRAPH`` instead."""
 
 
+class InvalidExecutionModeError(ValueError):
+    """Raised when an execution mode is not a recognized ExecutionMode value."""
+
+    def __init__(self, execution_mode: object) -> None:
+        super().__init__(f"Unknown execution_mode {execution_mode}. Expected 'graph' or 'eager'.")
+
+
 class OpQuantizerConfig(OpCompressionConfig[QuantizationSpec]):
     """
     Configuration class for quantization at the operation level.

src/coreai_opt/quantization/quantizer.py

@@ -30,6 +30,7 @@
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config.quantization_config import (
     ExecutionMode,
+    InvalidExecutionModeError,
     QATSchedule,
     QuantizerConfig,
 )
@@ -154,7 +155,7 @@ def __init__(
         elif execution_mode == ExecutionMode.EAGER:
             self._quantizer = _EagerQuantizer(model, config)
         else:
-            raise ValueError(f"Unsupported execution mode: {execution_mode}")
+            raise InvalidExecutionModeError(execution_mode)
 
         super().__init__(model, config)
 

tests/test_inspection.py

@@ -24,7 +24,7 @@
 from coreai_opt.inspection._eager_mode import _EagerOpDiscoveryMode
 from coreai_opt.inspection.types import BoundaryEdge, InputEdge, OpInfo
 from coreai_opt.palettization import KMeansPalettizer
-from coreai_opt.quantization import Quantizer
+from coreai_opt.quantization import InvalidExecutionModeError, Quantizer
 from coreai_opt.quantization.config.quantization_config import ExecutionMode
 
 execution_modes = pytest.mark.parametrize(
@@ -1138,9 +1138,9 @@ def test_eager_with_graph_module_raises_type_error(self) -> None:
             ModelInspector(gm, (torch.randn(1, 10),), execution_mode="eager")
 
     def test_invalid_execution_mode_raises(self) -> None:
-        """Verify ValueError for unrecognized execution mode."""
+        """Verify InvalidExecutionModeError for unrecognized execution mode."""
         model = nn.Linear(10, 5)
-        with pytest.raises(ValueError, match="Unknown execution_mode"):
+        with pytest.raises(InvalidExecutionModeError):
             ModelInspector(model, (torch.randn(1, 10),), execution_mode="invalid")
 
     def test_unsupported_compressor_raises(self) -> None: