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
@@ -174,8 +175,7 @@ def _compute_channel_mask(
             raise ValueError(
                 f"Invalid axis. Should be in range [{-weight.ndim}, {weight.ndim}), but got {axis}"
             )
-        if axis < 0:
-            axis += weight.ndim
+        axis = _normalize_axis(axis, weight.ndim)
 
         num_channels = weight.shape[axis]
         num_prune = math.floor(num_channels * sparsity)

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/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."""