[Bug]: ChannelStructured pruning crashes with RuntimeError on 1D tensors (PyTorch sum(dim=[]) scalar collapse)

[rohith500]

Summary

Applying ChannelStructured pruning to a 1D tensor parameter (e.g. 1D bias parameters, 1D normalization weights, or directly calling _MagnitudePruneImpl.compute_mask on a 1D tensor) crashes with:

RuntimeError: selected index k out of range

at torch.topk(channel_norms, num_keep, largest=True) in src/coreai_opt/pruning/spec/prune.py:192.

---

Context & Prior Art

This defect is closely related to PR #45 (fix(pruning): normalize a negative channel axis by @eyupcanakman, approved by @u-simha):

• In PR fix(pruning): normalize a negative channel axis #45, negative channel indexing (axis=-1) previously failed with this exact same error message (RuntimeError: selected index k out of range) because raw negative indices never matched d in range(ndim), causing reduce_dims to include all dimensions and collapsing channel_norms to a scalar.
• PR fix(pruning): normalize a negative channel axis #45 resolved negative indexing via _normalize_axis(axis, weight.ndim) and added tests for 4D (nn.Conv2d) and 2D (nn.Linear) tensors, as well as verifying normalize_axis(-1, 1) == 0.
• However, 1D tensors were never tested end-to-end with ChannelStructured.
• In addition, in open PR Add support for structured sparsity (block structure & n:m structure) #61 (Add support for structured sparsity by @u-simha), mask computation was abstracted into ChannelStructured.compute_mask_impl in src/coreai_opt/pruning/spec/scheme.py, inheriting the identical reduction logic.

---

Root Cause Analysis

In src/coreai_opt/pruning/spec/prune.py:188-192:

reduce_dims = [d for d in range(weight.ndim) if d != axis]
channel_norms = weight.abs().sum(dim=reduce_dims)

num_keep = num_channels - num_prune
_, keep_indices = torch.topk(channel_norms, num_keep, largest=True)

1. When weight.ndim == 1 and axis == 0 (or axis == -1, which normalizes to 0):

   reduce_dims = [d for d in range(1) if d != 0] # -> []

2. In PyTorch ATen C++ (ReduceOps.cpp), passing an empty dimension list dim=[] to sum(dim) does not act as an identity / no-op. Instead, PyTorch reduces across all dimensions into a 0-D scalar:

   >>> w = torch.tensor([1.0, 2.0, 3.0, 4.0])
   >>> w.sum(dim=[])
   tensor(10.)  # 0-D scalar with shape torch.Size([])

3. When torch.topk(channel_norms, num_keep) is invoked on a 0-D scalar tensor, PyTorch fails because the tensor has 0 dimensions:

   RuntimeError: selected index k out of range
   

Mathematical Invariant

In channel-structured pruning along axis, the $L_1$ norm of a channel slice is the sum over all non-channel dimensions:

$$ |W_{c}|_{1} = \sum_{d \neq \text{axis}} |W_{\dots, c, \dots}| $$

When $W$ is a 1D tensor (ndim == 1, axis == 0), each channel slice is simply an individual scalar element $w_i$. The $L_1$ norm of a scalar is its absolute value:

$$ |w_i|_{1} = |w_i| $$

Therefore, when reduce_dims is empty, no dimension reduction should occur; channel_norms is simply weight.abs().

---

Step-by-Step Reproduction

Minimal Reproduction Script

import torch
import torch.nn as nn
from coreai_opt.pruning import MagnitudePruner, MagnitudePrunerConfig
from coreai_opt.pruning.config import ModuleMagnitudePrunerConfig
from coreai_opt.pruning.spec import ChannelStructured, PruningSpec
from coreai_opt.pruning.spec.prune import _MagnitudePruneImpl

# Case 1: Direct mask computation on a 1D tensor
w = torch.tensor([1.0, 5.0, 2.0, 8.0])
_MagnitudePruneImpl._compute_channel_mask(w, sparsity=0.5, axis=0)
# -> RuntimeError: selected index k out of range

# Case 2: End-to-end model pruning on a 1D parameter (e.g. bias)
model = nn.Linear(4, 4, bias=True)
config = MagnitudePrunerConfig(
    global_config=ModuleMagnitudePrunerConfig(
        op_state_spec={
            "bias": PruningSpec(
                target_sparsity=0.5,
                pruning_scheme=ChannelStructured(axis=0),
            )
        }
    )
)
pruner = MagnitudePruner(model, config)
pruner.prepare((torch.randn(1, 4),))
# -> RuntimeError: selected index k out of range

Full Traceback

Traceback (most recent call last):
  File "<string>", line 20, in <module>
  File ".../src/coreai_opt/pruning/magnitude_pruner.py", line 124, in prepare
    prepared_model(*example_inputs)
  File ".../torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File ".../torch/nn/modules/linear.py", line 134, in forward
    return F.linear(input, self.weight, self.bias)
  File ".../torch/nn/utils/parametrize.py", line 440, in get_parametrized
    return parametrization()
  File ".../torch/nn/utils/parametrize.py", line 331, in forward
    x = self[0](self.original)
  File ".../src/coreai_opt/pruning/spec/prune.py", line 103, in forward
    new_mask = self.compute_mask(weight, self._sparsity, self._pruning_scheme)
  File ".../src/coreai_opt/pruning/spec/prune.py", line 149, in compute_mask
    return _MagnitudePruneImpl._compute_channel_mask(weight, sparsity, pruning_scheme.axis)
  File ".../src/coreai_opt/pruning/spec/prune.py", line 192, in _compute_channel_mask
    _, keep_indices = torch.topk(channel_norms, num_keep, largest=True)
RuntimeError: selected index k out of range

---

Proposed Permanent Canonical Fix

In src/coreai_opt/pruning/spec/prune.py:188-189 (and src/coreai_opt/pruning/spec/scheme.py in PR #61):

-        reduce_dims = [d for d in range(weight.ndim) if d != axis]
-        channel_norms = weight.abs().sum(dim=reduce_dims)
+        reduce_dims = [d for d in range(weight.ndim) if d != axis]
+        channel_norms = weight.abs().sum(dim=reduce_dims) if reduce_dims else weight.abs()

Why This Fix Is Complete and Sound:

1. Handles 1D Tensors: For 1D tensors, reduce_dims is empty, avoiding PyTorch's sum(dim=[]) scalar reduction and correctly setting channel_norms = weight.abs().
2. Preserves Multi-D Behavior: For 2D, 3D, 4D, etc., reduce_dims is non-empty, preserving identical existing behavior.
3. Axis Independence: Works seamlessly for axis=0 and axis=-1 after _normalize_axis.

---

Verification Plan

1. Unit Tests in tests/pruning/test_magnitude_pruner.py:
   • Add parametric tests for 1D tensor channel-structured pruning with positive and negative axes (axis=0, axis=-1).
   • Test end-to-end model preparation with 1D parameter pruning.
2. Quality Gates:
   • Run full make check (formatting, linting, docstrings, typing, tests) to ensure 100% compliance.

[u-simha]

@rohith500 Can you explain the use case here? A 1D tensor has no channel-structure; so even if we try to apply channel-structured purning on it, it degenerates to unstructured pruning.

I don't think supporting this makes sense; One potential fix can be to verify that channel structured pruning always receives a non-1D tensor.

[rohith500]

Hi @u-simha,

Thanks for the feedback!

Use Case Context:
Users hit this primarily through wildcard configs like op_state_spec={"*": ChannelStructured(axis=0)} on layers that have a 1D bias or normalization weights. Today, instead of an informative message, PyTorch crashes with an obscure ATen error (RuntimeError: selected index k out of range) because sum(dim=[]) collapses 1D tensors to a 0-D scalar before topk.

Agreement on Strict Validation:
You're completely right—a 1D tensor has no orthogonal channel structure, so channel pruning on it degenerates to element-wise (Unstructured) magnitude pruning. In fact, PyTorch core's own _validate_structured_pruning enforces this exact rule by requiring ndim >= 2.

Proposed Direction:
Your proposed fix is much cleaner and semantically sound. We can update PR #108 to validate weight.ndim >= 2:

if weight.ndim < 2:
    raise ValueError(
        f"Channel-structured pruning requires a tensor with at least 2 dimensions, "
        f"got shape {tuple(weight.shape)} with {weight.ndim} dims. "
        f"For 1D tensors, use Unstructured pruning instead."
    )

Would you like me to update PR #108 with this validation check and corresponding unit tests?

[u-simha]

Yes, the validation check makes sense.