From 1a9ee27cb30660239ba5efdac2cf450943b42b7c Mon Sep 17 00:00:00 2001
From: usimha <135899523+u-simha@users.noreply.github.com>
Date: Mon, 3 Aug 2026 16:52:55 -0700
Subject: [PATCH 1/3] Add block-structured and N:M-structured pruning schemes

Introduces BlockStructured and NMStructured PruningScheme
implementations alongside the existing Unstructured and
ChannelStructured, and refactors _MagnitudePruneImpl to delegate
mask computation to PruningScheme.compute_mask() instead of
switching on scheme type. _BlockSizeMismatchError moves to a
shared coreai_opt/_utils/errors.py so both quantization and
pruning specs can raise it.
---
 src/coreai_opt/_utils/errors.py               |  10 +
 src/coreai_opt/pruning/spec/__init__.py       |   4 +-
 src/coreai_opt/pruning/spec/errors.py         |  10 +
 src/coreai_opt/pruning/spec/prune.py          |  62 +-----
 src/coreai_opt/pruning/spec/scheme.py         | 197 +++++++++++++++++-
 src/coreai_opt/quantization/spec/errors.py    |   6 +-
 tests/fixtures/pruning.py                     |  18 +-
 tests/pruning/test_magnitude_pruner.py        | 183 +++++++++++++++-
 tests/pruning/test_pruning_config_and_spec.py |  82 +++++++-
 9 files changed, 504 insertions(+), 68 deletions(-)
 create mode 100644 src/coreai_opt/_utils/errors.py
 create mode 100644 src/coreai_opt/pruning/spec/errors.py

diff --git a/src/coreai_opt/_utils/errors.py b/src/coreai_opt/_utils/errors.py
new file mode 100644
index 0000000..23af82c
--- /dev/null
+++ b/src/coreai_opt/_utils/errors.py
@@ -0,0 +1,10 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Errors shared across compression domains (quantization, pruning, ...)."""
+
+
+class _BlockSizeMismatchError(ValueError):
+    """Raised when a tensor dimension is not divisible by the block size."""
diff --git a/src/coreai_opt/pruning/spec/__init__.py b/src/coreai_opt/pruning/spec/__init__.py
index 354345f..9e9df32 100644
--- a/src/coreai_opt/pruning/spec/__init__.py
+++ b/src/coreai_opt/pruning/spec/__init__.py
@@ -6,11 +6,13 @@
 """Pruning spec components: specs, schemes, and parametrizations."""
 
 from .prune import PruneImplBase, _MagnitudePruneImpl
-from .scheme import ChannelStructured, PruningScheme, Unstructured
+from .scheme import BlockStructured, ChannelStructured, NMStructured, PruningScheme, Unstructured
 from .spec import PruningSpec, default_weight_pruning_spec
 
 __all__ = [
+    "BlockStructured",
     "ChannelStructured",
+    "NMStructured",
     "PruneImplBase",
     "PruningScheme",
     "PruningSpec",
diff --git a/src/coreai_opt/pruning/spec/errors.py b/src/coreai_opt/pruning/spec/errors.py
new file mode 100644
index 0000000..700109f
--- /dev/null
+++ b/src/coreai_opt/pruning/spec/errors.py
@@ -0,0 +1,10 @@
+# Copyright 2026 Apple Inc.
+#
+# Use of this source code is governed by a BSD-3-Clause license that can
+# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+"""Errors raised by pruning scheme specifications."""
+
+from coreai_opt._utils.errors import _BlockSizeMismatchError
+
+__all__ = ["_BlockSizeMismatchError"]
diff --git a/src/coreai_opt/pruning/spec/prune.py b/src/coreai_opt/pruning/spec/prune.py
index 12c1e77..1918fce 100644
--- a/src/coreai_opt/pruning/spec/prune.py
+++ b/src/coreai_opt/pruning/spec/prune.py
@@ -7,7 +7,6 @@
 
 from __future__ import annotations
 
-import math
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any
 
@@ -19,7 +18,7 @@
 )
 from coreai_opt.config.spec import CompressionSimulatorBase
 
-from .scheme import ChannelStructured, PruningScheme
+from .scheme import PruningScheme
 
 if TYPE_CHECKING:
     # Imported only for type checking — runtime would be a circular import via
@@ -116,10 +115,11 @@ def with_args(cls, **kwargs: Any) -> _PartialConstructor[PruneImplBase]:
 
 @PruneImplBase.register("default")
 class _MagnitudePruneImpl(PruneImplBase):
-    """Magnitude-based pruning supporting unstructured and channel-structured schemes.
+    """Magnitude-based pruning that delegates to the configured pruning scheme.
 
     Prunes a given tensor to target sparsity by zero-ing out the smallest-magnitude
-    elements until desired sparsity is achieved.
+    elements, per whatever structural pattern ``pruning_scheme`` defines (unstructured,
+    channel-structured, block-structured, N:M-structured, or future schemes).
     """
 
     @staticmethod
@@ -128,7 +128,7 @@ def compute_mask(
         sparsity: float,
         pruning_scheme: PruningScheme,
     ) -> torch.Tensor:
-        """Compute a magnitude-based mask respecting the pruning scheme.
+        """Compute a magnitude-based mask by delegating to the pruning scheme.
 
         Args:
             weight (torch.Tensor): The weight tensor.
@@ -138,54 +138,4 @@ def compute_mask(
         Returns:
             torch.Tensor: Binary mask (1 = keep, 0 = prune).
         """
-        if sparsity == 0.0:
-            return torch.ones_like(weight)
-        if sparsity >= 1.0:
-            return torch.zeros_like(weight)
-
-        # TODO: Replace this with generic abstractions
-        if isinstance(pruning_scheme, ChannelStructured):
-            return _MagnitudePruneImpl._compute_channel_mask(weight, sparsity, pruning_scheme.axis)
-        return _MagnitudePruneImpl._compute_unstructured_mask(weight, sparsity)
-
-    @staticmethod
-    def _compute_unstructured_mask(weight: torch.Tensor, sparsity: float) -> torch.Tensor:
-        """Element-wise magnitude pruning."""
-        num_elements = weight.numel()
-        num_keep = num_elements - math.floor(num_elements * sparsity)
-        abs_weight = weight.abs()
-        _, topk_indices = torch.topk(abs_weight.flatten(), num_keep)
-        mask = torch.zeros(num_elements, dtype=weight.dtype, device=weight.device)
-        mask[topk_indices] = 1.0
-        return mask.reshape(weight.shape)
-
-    @staticmethod
-    def _compute_channel_mask(
-        weight: torch.Tensor,
-        sparsity: float,
-        axis: int,
-    ) -> torch.Tensor:
-        """Channel-structured magnitude pruning along *axis*.
-
-        Channel importance is measured by L1 norm. The least-important
-        channels are pruned entirely.
-        """
-        num_channels = weight.shape[axis]
-        num_prune = math.floor(num_channels * sparsity)
-
-        if num_prune == 0:
-            return torch.ones_like(weight)
-        if num_prune >= num_channels:
-            return torch.zeros_like(weight)
-
-        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)
-        channel_mask = torch.zeros(num_channels, dtype=weight.dtype, device=weight.device)
-        channel_mask[keep_indices] = 1.0
-
-        shape = [1] * weight.ndim
-        shape[axis] = num_channels
-        return channel_mask.view(shape).expand_as(weight)
+        return pruning_scheme.compute_mask(weight, sparsity)
diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py
index 7697619..9b873d4 100644
--- a/src/coreai_opt/pruning/spec/scheme.py
+++ b/src/coreai_opt/pruning/spec/scheme.py
@@ -7,18 +7,32 @@
 
 from __future__ import annotations
 
+import math
+from abc import abstractmethod
 from typing import Any, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, model_serializer
+import torch
+from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
 
 from coreai_opt._utils.registry_utils import ConfigRegistryMixin
 
+from .errors import _BlockSizeMismatchError
+
 
 class PruningScheme(BaseModel, ConfigRegistryMixin):
     """Base class for pruning scheme specifications.
 
     A pruning scheme defines the structural pattern of sparsity applied
-    to a tensor. Subclasses represent different ways of structuring the pruning.
+    to a tensor, and knows how to turn ``(weight, sparsity)`` into a binary
+    mask. Call the public :meth:`compute_mask` to get a mask; subclasses
+    implement the abstract :meth:`_compute_mask`, which handles sparsity
+    strictly between 0 and 1 (the 0.0 / 1.0 edge cases are handled once, in
+    the base class).
+
+    The sole exception is :class:`NMStructured`, whose achieved sparsity is
+    fixed by construction (``n / m``) rather than a free parameter — it
+    overrides :meth:`compute_mask` directly and ignores the ``sparsity``
+    argument entirely.
 
     Attributes:
         axis (int | None): The axis along which structured pruning is applied.
@@ -51,6 +65,32 @@ def _serialize_model(self) -> dict[str, Any]:
 
         return data
 
+    def compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        """Compute a binary pruning mask for the given weight tensor.
+
+        Args:
+            weight (torch.Tensor): The weight tensor to compute a mask for.
+            sparsity (float): Fraction of elements to prune, in [0, 1].
+
+        Returns:
+            torch.Tensor: Binary mask with the same shape as *weight* (1 = keep,
+            0 = prune).
+        """
+        if sparsity == 0.0:
+            return torch.ones_like(weight)
+        if sparsity >= 1.0:
+            return torch.zeros_like(weight)
+        return self._compute_mask(weight, sparsity)
+
+    @abstractmethod
+    def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        """Compute a mask for *sparsity* strictly between 0 and 1.
+
+        Subclasses implement the scheme-specific masking logic here; the
+        0.0 / 1.0 edge cases are already handled by :meth:`compute_mask`.
+        """
+        ...
+
 
 @PruningScheme.register("unstructured")
 class Unstructured(PruningScheme):
@@ -62,14 +102,163 @@ class Unstructured(PruningScheme):
 
     axis: Literal[None] = None
 
+    def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        num_elements = weight.numel()
+        num_keep = num_elements - math.floor(num_elements * sparsity)
+        abs_weight = weight.abs()
+        _, topk_indices = torch.topk(abs_weight.flatten(), num_keep)
+        mask = torch.zeros(num_elements, dtype=weight.dtype, device=weight.device)
+        mask[topk_indices] = 1.0
+        return mask.reshape(weight.shape)
+
 
 @PruningScheme.register("channel_structured")
 class ChannelStructured(PruningScheme):
     """Channel-structured pruning scheme.
 
     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).
+    Channel importance is determined by L1 norm of each channel.
     """
 
     axis: int = Field(default=0, description="Axis along which channels are pruned.")
+
+    def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        num_channels = weight.shape[self.axis]
+        num_prune = math.floor(num_channels * sparsity)
+
+        if num_prune == 0:
+            return torch.ones_like(weight)
+        if num_prune >= num_channels:
+            return torch.zeros_like(weight)
+
+        reduce_dims = [d for d in range(weight.ndim) if d != self.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)
+        channel_mask = torch.zeros(num_channels, dtype=weight.dtype, device=weight.device)
+        channel_mask[keep_indices] = 1.0
+
+        shape = [1] * weight.ndim
+        shape[self.axis] = num_channels
+        return channel_mask.view(shape).expand_as(weight)
+
+
+@PruningScheme.register("block_structured")
+class BlockStructured(PruningScheme):
+    """Block-structured pruning scheme.
+
+    Generalizes :class:`ChannelStructured` to prune contiguous blocks of
+    ``block_size`` slices along ``axis`` together, ranked by L2 norm
+    (``ChannelStructured`` is equivalent to ``block_size=1``, ranked by L1
+    norm — the two remain separate registered schemes).
+
+    Unlike Phoenix's reference implementation, a tensor whose size along
+    ``axis`` is not evenly divisible by ``block_size`` raises an error
+    instead of being padded.
+    """
+
+    axis: int = Field(default=0, description="Axis along which blocks are formed.")
+    block_size: int = Field(gt=0, description="Number of contiguous slices per block along axis.")
+
+    def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        num_along_axis = weight.shape[self.axis]
+        if num_along_axis % self.block_size != 0:
+            raise _BlockSizeMismatchError(
+                f"Tensor size {num_along_axis} along axis {self.axis} is not "
+                f"divisible by block_size {self.block_size}. Full tensor shape: "
+                f"{tuple(weight.shape)}"
+            )
+
+        num_blocks = num_along_axis // self.block_size
+        num_prune = math.floor(num_blocks * sparsity)
+
+        if num_prune == 0:
+            return torch.ones_like(weight)
+        if num_prune >= num_blocks:
+            return torch.zeros_like(weight)
+
+        # axis is now at position 0; all other dims keep their relative order.
+        moved = torch.movedim(weight, self.axis, 0)
+        other_dims = moved.shape[1:]
+        grouped = moved.view(num_blocks, self.block_size, *other_dims)
+        block_norms = grouped.pow(2).sum(dim=tuple(range(1, grouped.ndim))).sqrt()
+
+        num_keep = num_blocks - num_prune
+        _, keep_indices = torch.topk(block_norms, num_keep, largest=True)
+        block_mask = torch.zeros(num_blocks, dtype=weight.dtype, device=weight.device)
+        block_mask[keep_indices] = 1.0
+
+        mask = block_mask.repeat_interleave(self.block_size)
+        mask = mask.view(num_along_axis, *([1] * len(other_dims))).expand(
+            num_along_axis, *other_dims
+        )
+        # Restore axis to its original position.
+        return torch.movedim(mask, 0, self.axis)
+
+
+@PruningScheme.register("n_m_structured")
+class NMStructured(PruningScheme):
+    """N:M structured pruning scheme.
+
+    Zeroes exactly ``n`` smallest-magnitude elements out of every contiguous
+    group of ``m`` elements along ``axis`` — a hardware-friendly sparsity
+    pattern with a fixed sparsity ratio of ``n / m``.
+
+    Unlike other schemes, the achieved sparsity is fixed by construction and
+    does not depend on ``PruningSpec.target_sparsity``: :meth:`compute_mask`
+    overrides the base class directly and **ignores** its ``sparsity``
+    argument.
+
+    Unlike Phoenix's reference implementation, a tensor whose size along
+    ``axis`` is not evenly divisible by ``m`` raises an error instead of
+    being padded.
+    """
+
+    axis: int = Field(default=0, description="Axis along which N:M groups are formed.")
+    n: int = Field(ge=0, description="Number of smallest-magnitude elements zeroed per group of m.")
+    m: int = Field(gt=0, description="Group size along axis.")
+
+    @model_validator(mode="after")
+    def _validate_n_lt_m(self) -> NMStructured:
+        if self.n >= self.m:
+            raise ValueError(f"n ({self.n}) must be less than m ({self.m})")
+        return self
+
+    def _compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        """Unreachable: :meth:`compute_mask` is overridden directly below and never
+        delegates here. Defined only to satisfy the base class's abstract method.
+        """
+        raise NotImplementedError(
+            "NMStructured overrides compute_mask directly; _compute_mask is unused."
+        )
+
+    def compute_mask(self, weight: torch.Tensor, sparsity: float) -> torch.Tensor:
+        """Compute the N:M mask.
+
+        ``sparsity`` is accepted for interface compatibility with
+        :class:`PruneImplBase` but is ignored — achieved sparsity is fixed
+        at ``n / m`` by construction.
+        """
+        num_along_axis = weight.shape[self.axis]
+        if num_along_axis % self.m != 0:
+            raise _BlockSizeMismatchError(
+                f"Tensor size {num_along_axis} along axis {self.axis} is not "
+                f"divisible by m {self.m}. Full tensor shape: {tuple(weight.shape)}"
+            )
+
+        if self.n == 0:
+            return torch.ones_like(weight)
+
+        # axis is now at position -1; all other dims keep their relative order.
+        moved = torch.movedim(weight, self.axis, -1)
+        original_shape = moved.shape
+        grouped = moved.reshape(-1, self.m)
+
+        prune_idx = torch.argsort(grouped.abs(), dim=1, stable=True)[:, : self.n]
+        mask = torch.ones_like(grouped)
+        mask.scatter_(1, prune_idx, 0.0)
+
+        mask = mask.reshape(original_shape)
+        # Restore axis to its original position.
+        return torch.movedim(mask, -1, self.axis)
diff --git a/src/coreai_opt/quantization/spec/errors.py b/src/coreai_opt/quantization/spec/errors.py
index 2d537e3..8ba01af 100644
--- a/src/coreai_opt/quantization/spec/errors.py
+++ b/src/coreai_opt/quantization/spec/errors.py
@@ -3,6 +3,8 @@
 # Use of this source code is governed by a BSD-3-Clause license that can
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
+"""Errors raised by quantization spec components."""
 
-class _BlockSizeMismatchError(ValueError):
-    """Raised when a tensor dimension is not divisible by the block size."""
+from coreai_opt._utils.errors import _BlockSizeMismatchError
+
+__all__ = ["_BlockSizeMismatchError"]
diff --git a/tests/fixtures/pruning.py b/tests/fixtures/pruning.py
index f4392dc..454a856 100644
--- a/tests/fixtures/pruning.py
+++ b/tests/fixtures/pruning.py
@@ -11,7 +11,13 @@
 
 from coreai_opt import ExportBackend
 from coreai_opt.pruning import MagnitudePrunerConfig, ModuleMagnitudePrunerConfig, PruningSpec
-from coreai_opt.pruning.spec import ChannelStructured, PruningScheme, Unstructured
+from coreai_opt.pruning.spec import (
+    BlockStructured,
+    ChannelStructured,
+    NMStructured,
+    PruningScheme,
+    Unstructured,
+)
 
 
 @dataclass
@@ -21,7 +27,8 @@ class ParametrizedPruneConfigs:
     Attributes:
         config: MagnitudePrunerConfig instance.
         target_sparsity: Target sparsity fraction.
-        pruning_scheme: PruningScheme instance (Unstructured or ChannelStructured).
+        pruning_scheme: PruningScheme instance (Unstructured, ChannelStructured,
+            BlockStructured, or NMStructured).
         backend: Export backend (CoreML or CoreAI).
     """
 
@@ -53,7 +60,12 @@ def from_prune_params(
     params=[
         (target_sparsity, pruning_scheme, backend)
         for target_sparsity in [0.25, 0.5, 0.75]
-        for pruning_scheme in [Unstructured(), ChannelStructured(axis=0)]
+        for pruning_scheme in [
+            Unstructured(),
+            ChannelStructured(axis=0),
+            BlockStructured(axis=0, block_size=2),
+            NMStructured(axis=0, n=1, m=2),
+        ]
         for backend in [ExportBackend.CoreML, ExportBackend.CoreAI]
     ],
     ids=lambda p: f"sparsity:{p[0]}-scheme:{p[1].__class__.__name__}-backend:{p[2].value}",
diff --git a/tests/pruning/test_magnitude_pruner.py b/tests/pruning/test_magnitude_pruner.py
index 35aef7a..bf51d8e 100644
--- a/tests/pruning/test_magnitude_pruner.py
+++ b/tests/pruning/test_magnitude_pruner.py
@@ -23,7 +23,13 @@
     OpMagnitudePrunerConfig,
     PolynomialDecaySchedule,
 )
-from coreai_opt.pruning.spec import ChannelStructured, PruneImplBase, Unstructured
+from coreai_opt.pruning.spec import (
+    BlockStructured,
+    ChannelStructured,
+    NMStructured,
+    PruneImplBase,
+    Unstructured,
+)
 
 
 @pytest.fixture
@@ -486,6 +492,181 @@ def test_channel_structured_duplicate_norms(self) -> None:
             f"Expected exactly 3/5 channels pruned despite tied norms, got {num_pruned}"
         )
 
+    def test_block_structured_pruning_hand(self) -> None:
+        """Hand-written 8x4 tensor with block pruning (axis=0, block_size=2) at 50%.
+
+        Block magnitudes [5, 1, 4, 2] are interleaved so the two pruned blocks
+        (norms 1 and 2) are not contiguous — this keeps the case distinct from
+        plain channel-structured pruning of a contiguous leading/trailing slice.
+        """
+        model = nn.Linear(4, 8, bias=False)
+        with torch.no_grad():
+            model.weight.copy_(
+                torch.tensor(
+                    [
+                        [5.0, 5.0, 5.0, 5.0],
+                        [5.0, 5.0, 5.0, 5.0],
+                        [1.0, 1.0, 1.0, 1.0],
+                        [1.0, 1.0, 1.0, 1.0],
+                        [4.0, 4.0, 4.0, 4.0],
+                        [4.0, 4.0, 4.0, 4.0],
+                        [2.0, 2.0, 2.0, 2.0],
+                        [2.0, 2.0, 2.0, 2.0],
+                    ]
+                )
+            )
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.5,
+                        pruning_scheme=BlockStructured(axis=0, block_size=2),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 4),))
+
+        expected = torch.tensor(
+            [
+                [5.0, 5.0, 5.0, 5.0],
+                [5.0, 5.0, 5.0, 5.0],
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0],
+                [4.0, 4.0, 4.0, 4.0],
+                [4.0, 4.0, 4.0, 4.0],
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0],
+            ]
+        )
+        assert torch.equal(model.weight.detach(), expected)
+
+    def test_block_structured_non_leading_axis(self) -> None:
+        """Block pruning along axis=1 (input-feature dim) groups column pairs."""
+        model = nn.Linear(6, 4, bias=False)
+        with torch.no_grad():
+            model.weight.copy_(torch.tensor([[1.0, 1.0, 5.0, 5.0, 3.0, 3.0]]).expand(4, 6))
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.4,
+                        pruning_scheme=BlockStructured(axis=1, block_size=2),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 6),))
+
+        expected = torch.tensor([[0.0, 0.0, 5.0, 5.0, 3.0, 3.0]]).expand(4, 6)
+        assert torch.equal(model.weight.detach(), expected)
+
+    def test_block_structured_duplicate_norms(self) -> None:
+        """Blocks with tied L2 norms are still pruned to the exact target count.
+
+        Block magnitudes [1, 2, 2, 2, 5] (5 blocks of 2 rows each) at 60%
+        sparsity should prune exactly 3 of 5 blocks, despite the norm ties.
+        """
+        magnitudes = [1.0, 2.0, 2.0, 2.0, 5.0]
+        rows = [torch.full((4,), m) for m in magnitudes for _ in range(2)]
+        model = nn.Linear(4, 10, bias=False)
+        with torch.no_grad():
+            model.weight.copy_(torch.stack(rows))
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.6,
+                        pruning_scheme=BlockStructured(axis=0, block_size=2),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 4),))
+
+        weight = model.weight.detach()
+        num_pruned_blocks = sum(
+            1
+            for block_start in range(0, 10, 2)
+            if weight[block_start : block_start + 2].eq(0).all()
+        )
+        assert num_pruned_blocks == 3
+
+    def test_nm_structured_pruning_hand(self) -> None:
+        """NMStructured(n=2, m=4) enforces exact 2:4 sparsity per group, ignoring target_sparsity.
+
+        Each row is one group of 4 (axis=1): with strictly increasing values,
+        the 2 smallest elements per row are the exact ones zeroed.
+        """
+        model = nn.Linear(4, 8, bias=False)
+        original = torch.arange(1, 33, dtype=torch.float32).reshape(8, 4)
+        with torch.no_grad():
+            model.weight.copy_(original)
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.9,
+                        pruning_scheme=NMStructured(axis=1, n=2, m=4),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 4),))
+
+        expected = original.clone()
+        expected[:, :2] = 0.0
+        assert torch.equal(model.weight.detach(), expected)
+
+    def test_nm_structured_multi_group_per_axis_line(self) -> None:
+        """An 8-element axis with m=4, n=2 forms two independent groups per row."""
+        model = nn.Linear(8, 1, bias=False)
+        with torch.no_grad():
+            model.weight.copy_(torch.tensor([[4.0, 1.0, 3.0, 2.0, 8.0, 5.0, 7.0, 6.0]]))
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={"weight": PruningSpec(pruning_scheme=NMStructured(axis=1, n=2, m=4))}
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 8),))
+
+        # Group 1 = [4, 1, 3, 2]: two smallest (1, 2) zeroed -> keep 4, 3.
+        # Group 2 = [8, 5, 7, 6]: two smallest (5, 6) zeroed -> keep 8, 7.
+        expected = torch.tensor([[4.0, 0.0, 3.0, 0.0, 8.0, 0.0, 7.0, 0.0]])
+        assert torch.equal(model.weight.detach(), expected)
+
+    def test_nm_structured_n_zero_keeps_everything(self) -> None:
+        """NMStructured(n=0) is a degenerate no-op: the weight is left untouched."""
+        model = nn.Linear(8, 2, bias=False)
+        original = torch.arange(1, 17, dtype=torch.float32).reshape(2, 8)
+        with torch.no_grad():
+            model.weight.copy_(original)
+
+        config = MagnitudePrunerConfig(
+            global_config=ModuleMagnitudePrunerConfig(
+                op_state_spec={
+                    "weight": PruningSpec(
+                        target_sparsity=0.9,
+                        pruning_scheme=NMStructured(axis=1, n=0, m=4),
+                    )
+                }
+            )
+        )
+        pruner = MagnitudePruner(model, config)
+        pruner.prepare((torch.randn(1, 8),))
+
+        assert torch.equal(model.weight.detach(), original)
+
     def test_backprop_through_pruned_weights(self, linear_100x100_unique: nn.Linear) -> None:
         """Gradients flow through unpruned entries and are zero where the mask is zero."""
         pruner = MagnitudePruner(
diff --git a/tests/pruning/test_pruning_config_and_spec.py b/tests/pruning/test_pruning_config_and_spec.py
index 0eee33a..f41e6c8 100644
--- a/tests/pruning/test_pruning_config_and_spec.py
+++ b/tests/pruning/test_pruning_config_and_spec.py
@@ -15,7 +15,9 @@
     OpMagnitudePrunerConfig,
 )
 from coreai_opt.pruning.spec import (
+    BlockStructured,
     ChannelStructured,
+    NMStructured,
     PruneImplBase,
     PruningScheme,
     PruningSpec,
@@ -23,6 +25,7 @@
     _MagnitudePruneImpl,
     default_weight_pruning_spec,
 )
+from coreai_opt.pruning.spec.errors import _BlockSizeMismatchError
 
 
 class TestPruningSpec:
@@ -42,8 +45,17 @@ def test_default_spec(self) -> None:
             (0.75, Unstructured()),
             (0.5, ChannelStructured(axis=0)),
             (0.9, ChannelStructured(axis=1)),
+            (0.5, BlockStructured(axis=0, block_size=2)),
+            (0.9, NMStructured(axis=0, n=1, m=4)),
+        ],
+        ids=[
+            "25%-unstructured",
+            "75%-unstructured",
+            "50%-channel-ax0",
+            "90%-channel-ax1",
+            "50%-block-ax0",
+            "90%-nm-ax0",
         ],
-        ids=["25%-unstructured", "75%-unstructured", "50%-channel-ax0", "90%-channel-ax1"],
     )
     def test_custom_spec(self, target_sparsity: float, pruning_scheme: PruningScheme) -> None:
         """Custom spec values are accepted and stored correctly."""
@@ -93,6 +105,74 @@ def test_pruning_scheme_round_trip(
         assert isinstance(spec.pruning_scheme, expected_type)
         assert spec.pruning_scheme.axis == expected_axis
 
+    @pytest.mark.parametrize(
+        "scheme_dict,expected_axis,expected_block_size",
+        [
+            ({"type": "block_structured", "axis": 0, "block_size": 2}, 0, 2),
+            ({"type": "block_structured", "axis": 1, "block_size": 4}, 1, 4),
+        ],
+        ids=["block-ax0-size2", "block-ax1-size4"],
+    )
+    def test_block_structured_round_trip(
+        self, scheme_dict: dict, expected_axis: int, expected_block_size: int
+    ) -> None:
+        """BlockStructured constructed from dict resolves correctly."""
+        spec = PruningSpec(pruning_scheme=scheme_dict)
+        assert isinstance(spec.pruning_scheme, BlockStructured)
+        assert spec.pruning_scheme.axis == expected_axis
+        assert spec.pruning_scheme.block_size == expected_block_size
+
+    @pytest.mark.parametrize(
+        "scheme_dict,expected_axis,expected_n,expected_m",
+        [
+            ({"type": "n_m_structured", "axis": 0, "n": 1, "m": 4}, 0, 1, 4),
+            ({"type": "n_m_structured", "axis": 1, "n": 2, "m": 8}, 1, 2, 8),
+        ],
+        ids=["nm-ax0-1of4", "nm-ax1-2of8"],
+    )
+    def test_nm_structured_round_trip(
+        self, scheme_dict: dict, expected_axis: int, expected_n: int, expected_m: int
+    ) -> None:
+        """NMStructured constructed from dict resolves correctly."""
+        spec = PruningSpec(pruning_scheme=scheme_dict)
+        assert isinstance(spec.pruning_scheme, NMStructured)
+        assert spec.pruning_scheme.axis == expected_axis
+        assert spec.pruning_scheme.n == expected_n
+        assert spec.pruning_scheme.m == expected_m
+
+    def test_block_structured_invalid_block_size_raises(self) -> None:
+        """block_size <= 0 raises ValueError at construction time."""
+        with pytest.raises(ValueError):
+            BlockStructured(axis=0, block_size=0)
+
+    def test_block_structured_not_divisible_raises_via_spec(self) -> None:
+        """A tensor size not divisible by block_size raises _BlockSizeMismatchError."""
+        scheme = BlockStructured(axis=0, block_size=2)
+        with pytest.raises(_BlockSizeMismatchError):
+            scheme.compute_mask(torch.randn(5, 4), sparsity=0.5)
+
+    def test_nm_structured_not_divisible_raises_via_spec(self) -> None:
+        """A tensor size not divisible by m raises _BlockSizeMismatchError."""
+        scheme = NMStructured(axis=0, n=1, m=4)
+        with pytest.raises(_BlockSizeMismatchError):
+            scheme.compute_mask(torch.randn(5, 4), sparsity=0.0)
+
+    @pytest.mark.parametrize(
+        "n,m",
+        [(4, 4), (5, 4)],
+        ids=["n-equal-m", "n-greater-than-m"],
+    )
+    def test_nm_structured_invalid_n_m_raises_via_spec(self, n: int, m: int) -> None:
+        """n >= m raises ValueError, including when constructed via dict through the spec."""
+        with pytest.raises(ValueError, match="must be less than"):
+            PruningSpec(pruning_scheme={"type": "n_m_structured", "n": n, "m": m})
+
+    def test_nm_structured_n_zero_valid_construction(self) -> None:
+        """n=0 is a valid degenerate configuration regardless of m."""
+        scheme = NMStructured(n=0, m=4)
+        assert scheme.n == 0
+        assert scheme.m == 4
+
     def test_pruning_algo_round_trip(self) -> None:
         """Default pruning_algo string resolves to the class object."""
         spec = PruningSpec(pruning_algo="default")

From 4a5a1665f78ea808d68b31a6769a2bcc1fa2fb54 Mon Sep 17 00:00:00 2001
From: usimha <135899523+u-simha@users.noreply.github.com>
Date: Mon, 3 Aug 2026 17:33:13 -0700
Subject: [PATCH 2/3] Import _BlockSizeMismatchError directly from
 _utils.errors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Removes the pruning and quantization spec errors.py shims that
just re-exported _utils.errors._BlockSizeMismatchError, per review
feedback on apple/coreai-optimization#61 — all call sites now
import the shared error directly instead of going through a
per-domain copy.
---
 src/coreai_opt/pruning/spec/errors.py             | 10 ----------
 src/coreai_opt/pruning/spec/scheme.py             |  3 +--
 src/coreai_opt/quantization/spec/errors.py        | 10 ----------
 src/coreai_opt/quantization/spec/fake_quantize.py |  2 +-
 src/coreai_opt/quantization/spec/granularity.py   |  2 +-
 tests/pruning/test_pruning_config_and_spec.py     |  2 +-
 tests/quantization/test_quantization_spec.py      |  2 +-
 7 files changed, 5 insertions(+), 26 deletions(-)
 delete mode 100644 src/coreai_opt/pruning/spec/errors.py
 delete mode 100644 src/coreai_opt/quantization/spec/errors.py

diff --git a/src/coreai_opt/pruning/spec/errors.py b/src/coreai_opt/pruning/spec/errors.py
deleted file mode 100644
index 700109f..0000000
--- a/src/coreai_opt/pruning/spec/errors.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# Copyright 2026 Apple Inc.
-#
-# Use of this source code is governed by a BSD-3-Clause license that can
-# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
-
-"""Errors raised by pruning scheme specifications."""
-
-from coreai_opt._utils.errors import _BlockSizeMismatchError
-
-__all__ = ["_BlockSizeMismatchError"]
diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py
index 9b873d4..d898b22 100644
--- a/src/coreai_opt/pruning/spec/scheme.py
+++ b/src/coreai_opt/pruning/spec/scheme.py
@@ -14,10 +14,9 @@
 import torch
 from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
 
+from coreai_opt._utils.errors import _BlockSizeMismatchError
 from coreai_opt._utils.registry_utils import ConfigRegistryMixin
 
-from .errors import _BlockSizeMismatchError
-
 
 class PruningScheme(BaseModel, ConfigRegistryMixin):
     """Base class for pruning scheme specifications.
diff --git a/src/coreai_opt/quantization/spec/errors.py b/src/coreai_opt/quantization/spec/errors.py
deleted file mode 100644
index 8ba01af..0000000
--- a/src/coreai_opt/quantization/spec/errors.py
+++ /dev/null
@@ -1,10 +0,0 @@
-# Copyright 2026 Apple Inc.
-#
-# Use of this source code is governed by a BSD-3-Clause license that can
-# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
-
-"""Errors raised by quantization spec components."""
-
-from coreai_opt._utils.errors import _BlockSizeMismatchError
-
-__all__ = ["_BlockSizeMismatchError"]
diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py
index aafd476..8155527 100644
--- a/src/coreai_opt/quantization/spec/fake_quantize.py
+++ b/src/coreai_opt/quantization/spec/fake_quantize.py
@@ -15,6 +15,7 @@
 from torch.autograd import Function
 from torchao.quantization.pt2e import FakeQuantizeBase
 
+from coreai_opt._utils.errors import _BlockSizeMismatchError
 from coreai_opt._utils.spec_utils import (
     PartialConstructor as _PartialConstructor,
     with_args as _with_args,
@@ -27,7 +28,6 @@
 )
 from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor
 from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes
-from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 
 from .granularity import QuantizationGranularity
 from .qformulation import QuantizationFormulation
diff --git a/src/coreai_opt/quantization/spec/granularity.py b/src/coreai_opt/quantization/spec/granularity.py
index af0657a..f8a8d01 100644
--- a/src/coreai_opt/quantization/spec/granularity.py
+++ b/src/coreai_opt/quantization/spec/granularity.py
@@ -11,8 +11,8 @@
 import torch
 from pydantic import BaseModel, ConfigDict, Field, model_serializer
 
+from coreai_opt._utils.errors import _BlockSizeMismatchError
 from coreai_opt._utils.registry_utils import ConfigRegistryMixin
-from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 
 
 class QuantizationGranularity(BaseModel, ConfigRegistryMixin):
diff --git a/tests/pruning/test_pruning_config_and_spec.py b/tests/pruning/test_pruning_config_and_spec.py
index f41e6c8..285e8b7 100644
--- a/tests/pruning/test_pruning_config_and_spec.py
+++ b/tests/pruning/test_pruning_config_and_spec.py
@@ -9,6 +9,7 @@
 import torch
 import torch.nn as nn
 
+from coreai_opt._utils.errors import _BlockSizeMismatchError
 from coreai_opt.pruning.config import (
     MagnitudePrunerConfig,
     ModuleMagnitudePrunerConfig,
@@ -25,7 +26,6 @@
     _MagnitudePruneImpl,
     default_weight_pruning_spec,
 )
-from coreai_opt.pruning.spec.errors import _BlockSizeMismatchError
 
 
 class TestPruningSpec:
diff --git a/tests/quantization/test_quantization_spec.py b/tests/quantization/test_quantization_spec.py
index 8dc9efe..347b1d5 100644
--- a/tests/quantization/test_quantization_spec.py
+++ b/tests/quantization/test_quantization_spec.py
@@ -9,6 +9,7 @@
 from pydantic import ValidationError
 from torchao.quantization import MappingType as TorchAOMappingType
 
+from coreai_opt._utils.errors import _BlockSizeMismatchError
 from coreai_opt.quantization import QuantizationSpec
 from coreai_opt.quantization.spec import (
     PerBlockGranularity,
@@ -19,7 +20,6 @@
     default_activation_quantization_spec,
     default_weight_quantization_spec,
 )
-from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError
 from coreai_opt.quantization.spec.fake_quantize import (
     _DefaultFakeQuantizeImpl,
 )

From 94fc16fe0478d19337265193b4e61e69e51522b5 Mon Sep 17 00:00:00 2001
From: usimha <135899523+u-simha@users.noreply.github.com>
Date: Mon, 3 Aug 2026 17:47:25 -0700
Subject: [PATCH 3/3] Simplify divisibility docstrings in
 BlockStructured/NMStructured

State the axis-divisibility constraint directly instead of
contrasting with another implementation's padding behavior.
---
 src/coreai_opt/pruning/spec/scheme.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/coreai_opt/pruning/spec/scheme.py b/src/coreai_opt/pruning/spec/scheme.py
index d898b22..7d9b4a5 100644
--- a/src/coreai_opt/pruning/spec/scheme.py
+++ b/src/coreai_opt/pruning/spec/scheme.py
@@ -152,9 +152,8 @@ class BlockStructured(PruningScheme):
     (``ChannelStructured`` is equivalent to ``block_size=1``, ranked by L1
     norm — the two remain separate registered schemes).
 
-    Unlike Phoenix's reference implementation, a tensor whose size along
-    ``axis`` is not evenly divisible by ``block_size`` raises an error
-    instead of being padded.
+    The size of ``weight`` along ``axis`` must be evenly divisible by
+    ``block_size``.
     """
 
     axis: int = Field(default=0, description="Axis along which blocks are formed.")
@@ -209,9 +208,7 @@ class NMStructured(PruningScheme):
     overrides the base class directly and **ignores** its ``sparsity``
     argument.
 
-    Unlike Phoenix's reference implementation, a tensor whose size along
-    ``axis`` is not evenly divisible by ``m`` raises an error instead of
-    being padded.
+    The size of ``weight`` along ``axis`` must be evenly divisible by ``m``.
     """
 
     axis: int = Field(default=0, description="Axis along which N:M groups are formed.")