fix(pruning): normalize a negative channel axis (#45)

* fix(pruning): normalize a negative channel axis

ChannelStructured(axis=-1) pruned the wrong channels.
_compute_channel_mask compared each dim index against the raw axis when building its reduce list, so a negative axis excluded nothing and the per-channel L1 norms collapsed to a scalar.
Normalize the axis first, matching PerChannelGranularity, and reject an out-of-range axis with ValueError rather than letting it reach the reduction.

* refactor: move axis normalization into a shared util

`_compute_channel_mask` and `QuantizationGranularity._resolve_axis` both resolved a negative axis against the tensor rank. The shared helper lives in `torch_utils`. It resolves the axis and nothing else, so the range check stays where pruning already had it and the quantization path is unchanged.

Author: eyupcanakman

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/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.quantization.spec.errors import _BlockSizeMismatchError
 
 
@@ -104,9 +105,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/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."""