Merge branch 'main' into u/vineetgarg/block_activations

Author: vineet-g

src/coreai_opt/_utils/torch_utils.py

@@ -71,6 +71,23 @@ def flatten_tensors_to_list(obj: Any) -> list[torch.Tensor]:
     return []
 
 
+def normalize_axis(axis: int, ndim: int) -> int:
+    """Resolve a negative axis against a tensor rank.
+
+    A non-negative axis is returned unchanged. The caller must validate the range,
+    since an axis below ``-ndim`` stays negative here.
+
+    Args:
+        axis (int): Axis in standard Python style indexing.
+        ndim (int): Rank of the tensor the axis refers to.
+
+    Returns:
+        int: ``axis + ndim`` when axis is negative, otherwise axis unchanged.
+
+    """
+    return axis + ndim if axis < 0 else axis
+
+
 def get_module_name(model: torch.nn.Module, module: torch.nn.Module) -> str | None:
     """Find the fully qualified name of a module within a model.
 

src/coreai_opt/pruning/spec/prune.py

@@ -17,6 +17,7 @@
     PartialConstructor as _PartialConstructor,
     with_args as _with_args,
 )
+from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis
 from coreai_opt.config.spec import CompressionSimulatorBase
 
 from .scheme import ChannelStructured, PruningScheme
@@ -170,6 +171,12 @@ def _compute_channel_mask(
         Channel importance is measured by L1 norm. The least-important
         channels are pruned entirely.
         """
+        if not (-weight.ndim <= axis < weight.ndim):
+            raise ValueError(
+                f"Invalid axis. Should be in range [{-weight.ndim}, {weight.ndim}), but got {axis}"
+            )
+        axis = _normalize_axis(axis, weight.ndim)
+
         num_channels = weight.shape[axis]
         num_prune = math.floor(num_channels * sparsity)
 

src/coreai_opt/pruning/spec/scheme.py

@@ -70,6 +70,9 @@ class ChannelStructured(PruningScheme):
     Entire channels (slices along ``axis``) are pruned or kept together.
     Channel importance is determined by the pruning algorithm (e.g. L1 norm
     of each channel for magnitude-based pruning).
+
+    Note:
+        ``axis`` can be negatively indexed as per standard Python style indexing.
     """
 
     axis: int = Field(default=0, description="Axis along which channels are pruned.")

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -137,12 +137,11 @@ def _get_weight_input_names(
 
     Returns:
         Tuple of (module_name, param_name)
-        - module_name: e.g., "conv", "layer1.0"
+        - module_name: e.g., "conv", "layer1.0", or "" for a root-module parameter
         - param_name: e.g., "weight", "bias"
 
     Raises:
         ValueError: If node is not a weight quantization node
-        ValueError: If weight target path is invalid
 
     """
     if not _is_weight_fake_quant(fake_quant_node, module):
@@ -154,13 +153,10 @@ def _get_weight_input_names(
     # Extract module and parameter name from target path
     # e.g., "conv.weight" -> ("conv", "weight")
     # e.g., "layer1.0.weight" -> ("layer1.0", "weight")
+    # e.g., "weight" -> ("", "weight") for a parameter on the root module
     target_path = str(input_node.target)
     last_dot_idx = target_path.rfind(".")
-    if last_dot_idx == -1:
-        msg = f"Invalid weight target path: {target_path}"
-        raise ValueError(msg)
-
-    module_name = target_path[:last_dot_idx]
+    module_name = target_path[:last_dot_idx] if last_dot_idx != -1 else ""
     param_name = target_path[last_dot_idx + 1 :]
 
     return module_name, param_name

src/coreai_opt/quantization/spec/granularity.py

@@ -12,6 +12,7 @@
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
 from coreai_opt._utils.registry_utils import ConfigRegistryMixin as _ConfigRegistryMixin
+from coreai_opt._utils.torch_utils import normalize_axis as _normalize_axis
 from coreai_opt.config.spec import CompressionTargetTensor as _CompressionTargetTensor
 from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 
@@ -123,9 +124,7 @@ def _resolve_axis(granularity: QuantizationGranularity, tensor_ndim: int) -> int
         axis = granularity.axis
         if axis is None:
             return None
-        if axis < 0:
-            axis += tensor_ndim
-        return axis
+        return _normalize_axis(axis, tensor_ndim)
 
 
 @QuantizationGranularity.register("per_tensor")

tests/export/test_pt2e_weight_only_mil_export.py

@@ -83,3 +83,63 @@ def test_simple_model_weight_only_mil_export(
         export_backend=backend,
         prepared_model_output=prepared_model_output,
     )
+
+
+@pytest.mark.parametrize("dtype", ["int8", "uint8"])
+@pytest.mark.parametrize(
+    "granularity",
+    [
+        PerTensorGranularity(),
+        PerChannelGranularity(axis=0),
+        PerBlockGranularity(axis=0, block_size=2),
+    ],
+)
+def test_root_module_weight_only_mil_export(
+    dtype: str,
+    granularity: PerTensorGranularity | PerChannelGranularity | PerBlockGranularity,
+) -> None:
+    """Test weight-only quantization export to CoreML for a root-module weight.
+
+    A bare ``nn.Linear`` owns its weight directly, so the fake-quant input target is the
+    dot-less ``"weight"`` rather than ``"<submodule>.weight"``. The graph CoreML export
+    path rejected that with ``Invalid weight target path: weight`` instead of resolving it
+    to the root module.
+    """
+    model = torch.nn.Linear(8, 4)
+    model.eval()
+    model_input = torch.randn(2, 8)
+
+    config = QuantizerConfig(
+        global_config=ModuleQuantizerConfig(
+            op_state_spec={
+                "weight": QuantizationSpec(
+                    dtype=dtype,
+                    qscheme=QuantizationScheme.SYMMETRIC,
+                    granularity=granularity,
+                    fake_quantize_cls="default",
+                    qparam_calculator_cls="default",
+                    range_calculator_cls="minmax",
+                ),
+            },
+            op_input_spec=None,
+            op_output_spec=None,
+        ),
+    )
+
+    quantizer = Quantizer(model, config)
+    prepared_model = quantizer.prepare((model_input,))
+
+    with torch.no_grad():
+        prepared_model_output = prepared_model(model_input)
+
+    finalized_model = quantizer.finalize(backend=ExportBackend.CoreML)
+
+    export_utils.convert_and_verify(
+        finalized_model=finalized_model,
+        input_data=model_input,
+        expected_ops={
+            "constexpr_blockwise_shift_scale": 1,
+        },
+        export_backend=ExportBackend.CoreML,
+        prepared_model_output=prepared_model_output,
+    )

tests/pruning/test_magnitude_pruner.py

@@ -369,7 +369,8 @@ def test_channel_structured_pruning_hand(self) -> None:
         )
         assert torch.equal(model.weight.detach(), expected)
 
-    def test_channel_structured_conv2d(self) -> None:
+    @pytest.mark.parametrize("axis", [0, -4])
+    def test_channel_structured_conv2d(self, axis: int) -> None:
         """Channel-structured pruning on Conv2d zeros entire output filters."""
         torch.manual_seed(42)
         model = nn.Conv2d(3, 8, kernel_size=3, bias=False)
@@ -379,7 +380,7 @@ def test_channel_structured_conv2d(self) -> None:
                 op_state_spec={
                     "weight": PruningSpec(
                         target_sparsity=0.5,
-                        pruning_scheme=ChannelStructured(axis=0),
+                        pruning_scheme=ChannelStructured(axis=axis),
                     )
                 }
             )
@@ -395,6 +396,72 @@ def test_channel_structured_conv2d(self) -> None:
             filt = weight[i]
             assert filt.eq(0).all() or filt.ne(0).all(), f"Filter {i} is partially pruned"
 
+    @pytest.mark.parametrize("target_sparsity", [0.5, 0.75])
+    @pytest.mark.parametrize(
+        "negative_axis,positive_axis",
+        [(-1, 1), (-2, 0)],
+        ids=["last-dim", "first-dim"],
+    )
+    def test_channel_structured_negative_axis(
+        self, negative_axis: int, positive_axis: int, target_sparsity: float
+    ) -> None:
+        """A negative axis prunes the same channels as its positive equivalent."""
+
+        def prune_along(axis: int) -> torch.Tensor:
+            model = nn.Linear(4, 4, bias=False)
+            with torch.no_grad():
+                model.weight.copy_(
+                    torch.tensor(
+                        [
+                            [1.0, 2.0, 3.0, 4.0],
+                            [5.0, 6.0, 7.0, 8.0],
+                            [9.0, 10.0, 11.0, 12.0],
+                            [13.0, 14.0, 15.0, 16.0],
+                        ]
+                    )
+                )
+
+            config = MagnitudePrunerConfig(
+                global_config=ModuleMagnitudePrunerConfig(
+                    op_state_spec={
+                        "weight": PruningSpec(
+                            target_sparsity=target_sparsity,
+                            pruning_scheme=ChannelStructured(axis=axis),
+                        )
+                    }
+                )
+            )
+            pruner = MagnitudePruner(model, config)
+            pruner.prepare((torch.randn(1, 4),))
+            return model.weight.detach()
+
+        pruned = prune_along(negative_axis)
+        assert torch.equal(pruned, prune_along(positive_axis))
+
+        # L1 norms increase with index along both axes, so the highest indices survive.
+        num_keep = 4 - int(4 * target_sparsity)
+        kept = [i for i in range(4) if pruned.select(positive_axis, i).ne(0).any()]
+        assert kept == list(range(4 - num_keep, 4))
+
+    @pytest.mark.parametrize("axis", [-3, 2], ids=["below-range", "above-range"])
+    def test_channel_structured_axis_out_of_range(self, axis: int) -> None:
+        """An axis outside [-ndim, ndim) raises ValueError."""
+        model = nn.Linear(4, 4, bias=False)
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.5,
+                        pruning_scheme=ChannelStructured(axis=axis),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+
+        with pytest.raises(ValueError, match="Invalid axis"):
+            pruner.prepare((torch.randn(1, 4),))
+
     def test_linear_unstructured_conv2d_channel_structured(self) -> None:
         """Apply unstructured to Linear and channel-structured to Conv2d in same model."""
 

tests/test_utils/test_torch_utils.py

@@ -16,6 +16,7 @@
     mmap_module_state_dict,
     move_model_to_eval,
     move_model_to_train,
+    normalize_axis,
 )
 
 
@@ -93,6 +94,29 @@ def test_normalize_module_fqn(raw: str, expected: str) -> None:
         assert normalize_module_fqn(raw) == expected
 
 
+class TestNormalizeAxis:
+    """Test normalize_axis resolution of negative axes."""
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        ("axis", "ndim", "expected"),
+        [
+            (0, 2, 0),
+            (1, 2, 1),
+            (-1, 2, 1),
+            (-2, 2, 0),
+            (-4, 4, 0),
+            (-1, 1, 0),
+            # Out of range passes through unchanged; the caller validates.
+            (-5, 2, -3),
+            (3, 2, 3),
+        ],
+    )
+    def test_resolves_to_non_negative(axis: int, ndim: int, expected: int) -> None:
+        """A negative axis resolves to its non-negative equivalent."""
+        assert normalize_axis(axis, ndim) == expected
+
+
 class TestMmapModuleStateDict:
     """Test mmap_module_state_dict serialization and reload."""