refactor: use dataclass for weight axis defaults table (#30)

* refactor: use dataclass for weight axis defaults table

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

* address comments

* import privately

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

---------

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>

Author: pkmandke

src/coreai_opt/_utils/torch_utils.py

@@ -33,6 +33,19 @@
 E8M0_EXPONENT_BIAS: Final[int] = 127
 F32_MIN_NORMAL: Final[float] = 2**-126  # ~1.175e-38
 
+# Maps a graph-mode aten ``OpOverload`` to the ``nn.Module`` type it corresponds
+# to.
+ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[torch.nn.Module]] = {
+    torch.ops.aten.conv1d.default: torch.nn.Conv1d,
+    torch.ops.aten.conv2d.default: torch.nn.Conv2d,
+    torch.ops.aten.conv3d.default: torch.nn.Conv3d,
+    torch.ops.aten.conv_transpose1d.default: torch.nn.ConvTranspose1d,
+    torch.ops.aten.conv_transpose2d.input: torch.nn.ConvTranspose2d,
+    torch.ops.aten.conv_transpose3d.input: torch.nn.ConvTranspose3d,
+    torch.ops.aten.linear.default: torch.nn.Linear,
+    torch.ops.aten.embedding.default: torch.nn.Embedding,
+}
+
 
 class NamedModule(NamedTuple):
     """NamedTuple for holding name and module info"""

src/coreai_opt/quantization/_axis_defaults.py

@@ -13,12 +13,16 @@
 
 import logging
 from collections import defaultdict
+from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
 import torch.nn.utils.parametrize as P
 from torch.fx import GraphModule
 
+from coreai_opt._utils.torch_utils import (
+    ATEN_OP_TO_MODULE_TYPE as _ATEN_OP_TO_MODULE_TYPE,
+)
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
 from coreai_opt.quantization.spec.granularity import (
@@ -35,43 +39,38 @@
 _WeightFQMap = defaultdict[FakeQuantizeImplBase, list[_ConsumerInfo]]
 
 
-# Per-channel defaults: quantize along the OUTPUT channel axis.
-# Conv weights are [out_ch, in_ch, ...]; ConvTranspose weights are [in_ch, out_ch, ...].
-# Embedding weights are [num_embeddings, embedding_dim].
-_PER_CHANNEL_WEIGHT_AXIS_DEFAULTS: dict[type[nn.Module], int] = {
-    nn.Conv1d: 0,
-    nn.Conv2d: 0,
-    nn.Conv3d: 0,
-    nn.ConvTranspose1d: 1,
-    nn.ConvTranspose2d: 1,
-    nn.ConvTranspose3d: 1,
-    nn.Linear: 0,
-    nn.Embedding: 0,
-}
+@dataclass(frozen=True)
+class _WeightAxisSpec:
+    """Default axis indices for a single ``nn.Module`` type.
 
-# Per-block defaults: quantize along the INPUT channel axis (reduction dimension)
-# Embedding weights are [num_embeddings, embedding_dim].
-_PER_BLOCK_WEIGHT_AXIS_DEFAULTS: dict[type[nn.Module], int] = {
-    nn.Conv1d: 1,
-    nn.Conv2d: 1,
-    nn.Conv3d: 1,
-    nn.ConvTranspose1d: 0,
-    nn.ConvTranspose2d: 0,
-    nn.ConvTranspose3d: 0,
-    nn.Linear: 1,
-    nn.Embedding: 1,
-}
+    Attributes:
+        per_channel_axis (int): Output-channel axis for per-channel granularity.
+        per_block_axis (int): Input-channel (reduction) axis for per-block granularity.
+    """
+
+    per_channel_axis: int
+    per_block_axis: int
 
-# Maps aten OpOverload -> nn.Module type for graph-mode op to module type resolution.
-_ATEN_OP_TO_MODULE_TYPE: dict[torch._ops.OpOverload, type[nn.Module]] = {
-    torch.ops.aten.conv1d.default: nn.Conv1d,
-    torch.ops.aten.conv2d.default: nn.Conv2d,
-    torch.ops.aten.conv3d.default: nn.Conv3d,
-    torch.ops.aten.conv_transpose1d.default: nn.ConvTranspose1d,
-    torch.ops.aten.conv_transpose2d.input: nn.ConvTranspose2d,
-    torch.ops.aten.conv_transpose3d.input: nn.ConvTranspose3d,
-    torch.ops.aten.linear.default: nn.Linear,
-    torch.ops.aten.embedding.default: nn.Embedding,
+    def default_axis_for(self, granularity: QuantizationGranularity) -> int | None:
+        """Return the default axis for the given granularity, or ``None`` if not applicable."""
+        if isinstance(granularity, PerChannelGranularity):
+            return self.per_channel_axis
+        if isinstance(granularity, PerBlockGranularity):
+            return self.per_block_axis
+        return None
+
+
+# Conv weights are [out_ch, in_ch, ...]; ConvTranspose weights are [in_ch, out_ch, ...].
+# Embedding weights are [num_embeddings, embedding_dim].
+_WEIGHT_AXIS_SPECS: dict[type[nn.Module], _WeightAxisSpec] = {
+    nn.Conv1d: _WeightAxisSpec(0, 1),
+    nn.Conv2d: _WeightAxisSpec(0, 1),
+    nn.Conv3d: _WeightAxisSpec(0, 1),
+    nn.ConvTranspose1d: _WeightAxisSpec(1, 0),
+    nn.ConvTranspose2d: _WeightAxisSpec(1, 0),
+    nn.ConvTranspose3d: _WeightAxisSpec(1, 0),
+    nn.Linear: _WeightAxisSpec(0, 1),
+    nn.Embedding: _WeightAxisSpec(0, 1),
 }
 
 
@@ -220,12 +219,11 @@ def _collect_weight_fq_entries_eager(model: nn.Module) -> _WeightFQMap:
 def _apply_defaults(fq_map: _WeightFQMap) -> None:
     """Apply default weight axes and raise on unresolved or conflicting entries.
 
-    For each FQ in the map, selects the per-channel or per-block defaults
-    table based on the granularity type. Resolves each consumer's identifier
-    to a module type (aten ``OpOverload`` entries from graph-mode are mapped via
-    ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis. When a
-    single FQ has multiple consumers (shared weight), all consumers must
-    agree on the same default axis.
+    For each FQ in the map, resolves each consumer's identifier to a module
+    type (aten ``OpOverload`` entries from graph-mode are mapped via
+    ``_ATEN_OP_TO_MODULE_TYPE``), then looks up the default axis from
+    ``_WEIGHT_AXIS_SPECS``. When a single FQ has multiple consumers (shared
+    weight), all consumers must agree on the same default axis.
 
     Args:
         fq_map (_WeightFQMap): Map from fake-quantize instance to its
@@ -243,14 +241,6 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None:
         if not _granularity_needs_axis_default(fq.granularity):
             continue
 
-        # get appropriate defaults table
-        if isinstance(fq.granularity, PerChannelGranularity):
-            axis_defaults = _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS
-        elif isinstance(fq.granularity, PerBlockGranularity):
-            axis_defaults = _PER_BLOCK_WEIGHT_AXIS_DEFAULTS
-        else:
-            continue
-
         # Resolve each consumer to a module type and look up its default axis
         resolved_axes: set[int] = set()
         all_consumer_names: list[str] = []
@@ -261,8 +251,8 @@ def _apply_defaults(fq_map: _WeightFQMap) -> None:
             else:
                 module_type = module_type_or_op
 
-            # pick a default based on this consumer
-            default_axis = axis_defaults.get(module_type) if module_type is not None else None
+            spec = _WEIGHT_AXIS_SPECS.get(module_type) if module_type is not None else None
+            default_axis = spec.default_axis_for(fq.granularity) if spec is not None else None
             if default_axis is not None:
                 resolved_axes.add(default_axis)
 

tests/quantization/test_axis_defaults.py

@@ -19,9 +19,9 @@
     QuantizerConfig,
 )
 from coreai_opt.quantization._axis_defaults import (
-    _PER_BLOCK_WEIGHT_AXIS_DEFAULTS,
-    _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS,
+    _WEIGHT_AXIS_SPECS,
     _apply_defaults,
+    _WeightAxisSpec,
     _WeightFQMap,
 )
 from coreai_opt.quantization.spec import default_activation_quantization_spec
@@ -184,9 +184,9 @@ def test_axis_none_resolved(
         prepared = Quantizer(model, config).prepare((make_input(),))
 
         if isinstance(granularity, PerChannelGranularity):
-            expected_axis = _PER_CHANNEL_WEIGHT_AXIS_DEFAULTS[module_type]
+            expected_axis = _WEIGHT_AXIS_SPECS[module_type].per_channel_axis
         else:
-            expected_axis = _PER_BLOCK_WEIGHT_AXIS_DEFAULTS[module_type]
+            expected_axis = _WEIGHT_AXIS_SPECS[module_type].per_block_axis
 
         weight_fqs = _get_weight_fqs(prepared)
         assert len(weight_fqs) == 1
@@ -293,6 +293,25 @@ def _activation_only_config(
     )
 
 
+class TestWeightAxisSpec:
+    """Verify _WeightAxisSpec.default_axis_for maps granularity types to axes."""
+
+    def test_per_channel_returns_per_channel_axis(self):
+        """PerChannelGranularity resolves to the spec's per_channel_axis."""
+        spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1)
+        assert spec.default_axis_for(PerChannelGranularity(axis=None)) == 0
+
+    def test_per_block_returns_per_block_axis(self):
+        """PerBlockGranularity resolves to the spec's per_block_axis."""
+        spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1)
+        assert spec.default_axis_for(PerBlockGranularity(axis=None, block_size=2)) == 1
+
+    def test_other_granularity_returns_none(self):
+        """A granularity that is neither per-channel nor per-block returns None."""
+        spec = _WeightAxisSpec(per_channel_axis=0, per_block_axis=1)
+        assert spec.default_axis_for(PerTensorGranularity()) is None
+
+
 class TestActivationAxisValidation:
     """Verify that activation FQs with unresolved axis=None are caught at prepare time."""