From d382a9ea1050a815cc0a82af5ffafdbbda9808ac Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Fri, 7 Aug 2026 11:32:52 -0700
Subject: [PATCH 01/10] enhanced logging for incompatible layers with
 quant/palettization blockwise

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 .../config/spec/compression_simulator.py      | 46 ++++++++++
 src/coreai_opt/palettization/_source_names.py | 43 +++++++++
 .../palettization/kmeans/palettizer.py        | 16 ++--
 .../palettization/spec/fake_palettize.py      | 29 +++++--
 .../quantization/_eager/quantizer.py          |  7 ++
 .../quantization/_graph/quantizer.py          |  7 ++
 src/coreai_opt/quantization/_source_names.py  | 87 +++++++++++++++++++
 .../quantization/spec/fake_quantize.py        | 26 ++++--
 tests/palettization/test_kmeans_palettizer.py | 32 +++++++
 tests/quantization/test_eager_quant.py        | 37 ++++++++
 .../quantization/test_graph_mode_quantizer.py | 36 ++++++++
 11 files changed, 346 insertions(+), 20 deletions(-)
 create mode 100644 src/coreai_opt/palettization/_source_names.py
 create mode 100644 src/coreai_opt/quantization/_source_names.py

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index d3324e23..2f4ac427 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -26,6 +26,52 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     compression simulation is performed during training.
     """
 
+    # Model-level identity of the tensor this simulator compresses, recorded
+    # during prepare() so that forward-time diagnostics can name the tensor.
+    # These are class-level defaults rather than __init__ assignments so that
+    # subclass constructor chains stay untouched. Being plain strings, they are
+    # kept out of state_dict by nn.Module.__setattr__.
+    _source_module_name: str | None = None
+    _source_param_name: str | None = None
+
+    def set_source_name(self, module_name: str, param_name: str | None = None) -> None:
+        """Record which model tensor this simulator compresses.
+
+        A simulator cannot discover its own position in the model from inside
+        ``forward``, so callers record it during ``prepare()`` instead. When one
+        simulator is shared by several modules, the first name recorded wins, so
+        that the reported name is stable across runs.
+
+        Args:
+            module_name: Fully-qualified name of the owning module as it appears
+                in ``named_modules()``, or ``""`` for a root-module parameter.
+            param_name: Local parameter name (e.g. ``"weight"``). ``None`` when
+                the simulator does not act on a named parameter.
+        """
+        if self._source_module_name is not None or self._source_param_name is not None:
+            return
+        self._source_module_name = module_name
+        self._source_param_name = param_name
+
+    @property
+    def source_name(self) -> str:
+        """FQN of the compressed tensor (e.g. ``"layers.0.q_proj.weight"``).
+
+        Returns ``"<unknown>"`` when no name was recorded, which is the case for
+        activations and for any simulator created outside a ``prepare()`` pass.
+        """
+        parts = [part for part in (self._source_module_name, self._source_param_name) if part]
+        return ".".join(parts) if parts else "<unknown>"
+
+    @property
+    def source_module_name(self) -> str | None:
+        """FQN of the owning module, suitable as a ``module_name_configs`` key.
+
+        ``None`` when no name was recorded or when the tensor belongs to the
+        root module (which cannot be addressed by module name).
+        """
+        return self._source_module_name or None
+
     @abstractmethod
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """
diff --git a/src/coreai_opt/palettization/_source_names.py b/src/coreai_opt/palettization/_source_names.py
new file mode 100644
index 00000000..914033e1
--- /dev/null
+++ b/src/coreai_opt/palettization/_source_names.py
@@ -0,0 +1,43 @@
+# 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
+
+"""Record which model tensor each fake-palettize module palettizes.
+
+``_FakePalettizeImplBase`` warns and disables itself from inside ``forward``
+when a tensor is incompatible with the configured granularity or cluster
+dimension, but a module cannot discover its own position in the model at that
+point. The pass here runs during ``prepare()``, before centroids are computed,
+and stamps each fake-palettize module with the FQN of the parameter it
+palettizes so the warning can name the offending weight.
+
+Palettization targets weights only and is eager-only, so a single walk over
+``named_modules()`` covers every case.
+"""
+
+from __future__ import annotations
+
+import torch.nn as nn
+import torch.nn.utils.parametrize as P
+
+from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase
+
+
+def record_weight_source_names(model: nn.Module) -> None:
+    """Stamp fake-palettize modules with the FQN of the parameter they palettize.
+
+    Fake-palettize modules live in the ``ParametrizationList`` registered for the
+    parameter they compress, so the owning module name and the parameter name
+    both come straight from ``named_modules()``.
+
+    Args:
+        model (nn.Module): The prepared model.
+    """
+    for module_name, module in model.named_modules(remove_duplicate=True):
+        if not P.is_parametrized(module):
+            continue
+        for param_name, parametrizations in module.parametrizations.items():
+            for fake_palett in parametrizations:
+                if isinstance(fake_palett, _FakePalettizeImplBase):
+                    fake_palett.set_source_name(module_name, param_name)
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index c42161fe..c66676c8 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -34,6 +34,9 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
+from coreai_opt.palettization._source_names import (
+    record_weight_source_names as _record_weight_source_names,
+)
 from coreai_opt.palettization.base_palettizer import _BasePalettizer
 from coreai_opt.palettization.config.palettization_config import (
     KMeansPalettizerConfig,
@@ -100,9 +103,6 @@ def _calculate_centroids_for_module(
     except Exception as e:
         raise RuntimeError(f"Centroid calculation failed for layer {layer_name!r}") from e
 
-    if fp_module._disabled:
-        fp_module._disabled_reason = f"layer {layer_name!r}"
-
     return fp_module
 
 
@@ -198,6 +198,11 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
+        # Record each fake-palettize module's parameter FQN so that the
+        # incompatibility warnings raised while computing centroids below can
+        # name the offending weight.
+        _record_weight_source_names(prepared_model)
+
         # Save example inputs for later use in calibration
         self._example_inputs = tuple([ip.detach().clone() for ip in example_inputs])
 
@@ -518,10 +523,7 @@ def _calculate_centroids_parallel(self, num_workers: int) -> None:
 
         for info, new_fp in zip(fp_info, results, strict=True):
             if getattr(new_fp, "_disabled", False):
-                logger.warning(
-                    f"Disabling palettization for a module: "
-                    f"{getattr(new_fp, '_disabled_reason', '')}"
-                )
+                logger.warning("Disabling palettization for weight '%s'", new_fp.source_name)
             # ParametrizationList supports item assignment; this swaps the
             # worker's mutated module into the live model without touching
             # the surrounding parametrization registration.
diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py
index ee4f7e36..c82ac31f 100644
--- a/src/coreai_opt/palettization/spec/fake_palettize.py
+++ b/src/coreai_opt/palettization/spec/fake_palettize.py
@@ -78,6 +78,25 @@ def is_disabled(self) -> bool:
         """Return True if fake palettization has been disabled."""
         return self._disabled
 
+    def _warn_and_disable(self, setting: str, error: Exception, shape: torch.Size) -> None:
+        """Log a warning naming the offending weight and disable this module.
+
+        Names the tensor and its shape so that the palettization config can be
+        corrected without having to hunt for the layer first. The name is
+        recorded during ``prepare()``; it falls back to ``"<unknown>"`` for
+        modules created outside that path.
+        """
+        logger.warning(
+            "Tensor '%s' (shape: %s) incompatible with %s: %s. Skipping palettization.",
+            self.source_name,
+            tuple(shape),
+            setting,
+            # The granularity and cluster_dim messages end in a period; drop it
+            # so the sentence this builds has exactly one.
+            str(error).rstrip("."),
+        )
+        self._disabled = True
+
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """Apply fake palettization to input tensor"""
         # If permanently disabled due to incompatibility, return original tensor
@@ -89,16 +108,10 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
             try:
                 lut, indices = self._calculate_centroids(tensor)
             except _IncompatibleGranularityError as e:
-                logger.warning(
-                    f"Tensor incompatible with granularity: {e}. Skipping palettization."
-                )
-                self._disabled = True
+                self._warn_and_disable("granularity", e, tensor.shape)
                 return tensor
             except _IncompatibleClusterDimError as e:
-                logger.warning(
-                    f"Tensor incompatible with cluster_dim: {e}. Skipping palettization."
-                )
-                self._disabled = True
+                self._warn_and_disable("cluster_dim", e, tensor.shape)
                 return tensor
 
             self.lut = lut.detach()
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 7f00a7f2..35d49458 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -39,6 +39,9 @@
     disable_activation_fake_quant,
     enable_weight_fake_quant,
 )
+from coreai_opt.quantization._source_names import (
+    record_weight_source_names_eager as _record_weight_source_names,
+)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     ModuleQuantizerConfig,
@@ -223,6 +226,10 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
+        # Record each weight FQ's parameter FQN so that the block-size warning
+        # raised during the upcoming forward pass can name the offending weight.
+        _record_weight_source_names(model)
+
         _apply_weight_axis_defaults(model)
         validate_activation_axes(model)
 
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 09e0a484..b259505a 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -54,6 +54,9 @@
     disable_activation_fake_quant,
     enable_weight_fake_quant,
 )
+from coreai_opt.quantization._source_names import (
+    record_weight_source_names_graph as _record_weight_source_names,
+)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     KVCacheQuantConfig,
@@ -1298,6 +1301,10 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
+        # Record each weight FQ's parameter FQN so that the block-size warning
+        # raised during the upcoming forward pass can name the offending weight.
+        _record_weight_source_names(model)
+
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)
 
diff --git a/src/coreai_opt/quantization/_source_names.py b/src/coreai_opt/quantization/_source_names.py
new file mode 100644
index 00000000..1e561959
--- /dev/null
+++ b/src/coreai_opt/quantization/_source_names.py
@@ -0,0 +1,87 @@
+# 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
+
+"""Record which model tensor each weight fake-quantize module quantizes.
+
+``FakeQuantizeImplBase`` warns and disables itself from inside ``forward`` when a
+tensor is incompatible with the configured block size, but a module cannot
+discover its own position in the model at that point. The passes here run during
+``prepare()``, before the first forward pass, and stamp each weight
+fake-quantize module with the FQN of the parameter it quantizes so the warning
+can name the offending weight.
+
+Weights only: activation fake-quantize modules have no backing parameter, so
+they keep the unnamed warning.
+
+Like :mod:`coreai_opt.quantization._axis_defaults`, this module deliberately
+implements its own graph and eager walks instead of reusing the helpers in
+``_graph``/``_eager``. Both mode-specific quantizers import it, so importing
+either subpackage from here would create a circular import.
+"""
+
+from __future__ import annotations
+
+import torch.nn as nn
+import torch.nn.utils.parametrize as P
+from torch.fx import GraphModule
+
+from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+
+
+def record_weight_source_names_graph(model: GraphModule) -> None:
+    """Stamp weight fake-quantize modules in a graph-mode ``GraphModule``.
+
+    A weight fake-quantize node takes its value from the ``get_attr`` node for
+    the parameter, whose target is the dotted parameter FQN (e.g.
+    ``"layer1.0.weight"``), so the name needs no inference. This mirrors
+    ``_graph._prepare_for_export._get_weight_input_names``.
+
+    Args:
+        model (GraphModule): The prepared graph-mode ``GraphModule``.
+    """
+    modules = dict(model.named_modules(remove_duplicate=False))
+
+    for node in model.graph.nodes:
+        if node.op != "call_module":
+            continue
+        fake_quant = modules.get(str(node.target))
+        if not isinstance(fake_quant, FakeQuantizeImplBase):
+            continue
+        if fake_quant.quantization_target != CompressionTargetTensor.WEIGHT:
+            continue
+
+        # Activation fake-quantize nodes read from an op rather than a parameter.
+        input_node = node.args[0]
+        if input_node.op != "get_attr":
+            continue
+
+        # "layer1.0.weight" -> ("layer1.0", "weight"); a root-module parameter
+        # such as "weight" has no module part.
+        target_path = str(input_node.target)
+        module_name, _, param_name = target_path.rpartition(".")
+        fake_quant.set_source_name(module_name, param_name)
+
+
+def record_weight_source_names_eager(model: nn.Module) -> None:
+    """Stamp weight fake-quantize modules in an eager-mode model.
+
+    Weight fake-quantize modules live in the ``ParametrizationList`` registered
+    for the parameter they quantize, so the owning module name and the parameter
+    name both come straight from ``named_modules()``.
+
+    Args:
+        model (nn.Module): The prepared eager-mode model.
+    """
+    for module_name, module in model.named_modules(remove_duplicate=True):
+        if not P.is_parametrized(module):
+            continue
+        for param_name, parametrizations in module.parametrizations.items():
+            for fake_quant in parametrizations:
+                if not isinstance(fake_quant, FakeQuantizeImplBase):
+                    continue
+                if fake_quant.quantization_target != CompressionTargetTensor.WEIGHT:
+                    continue
+                fake_quant.set_source_name(module_name, param_name)
diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py
index 9ff72aa9..8604176b 100644
--- a/src/coreai_opt/quantization/spec/fake_quantize.py
+++ b/src/coreai_opt/quantization/spec/fake_quantize.py
@@ -128,13 +128,29 @@ def enable_observer(self, enabled: bool = True) -> None:
             return
         super().enable_observer(enabled)
 
-    def _warn_and_disable(self, error: _BlockSizeMismatchError) -> None:
-        """Log a warning and permanently disable this module."""
+    def _warn_and_disable(self, error: _BlockSizeMismatchError, shape: torch.Size) -> None:
+        """Log a warning and permanently disable this module.
+
+        Names the offending tensor and its shape so that the granularity can be
+        corrected without having to hunt for the layer first. The name is only
+        available for weights (recorded during ``prepare()``); activations fall
+        back to ``"<unknown>"``.
+        """
+        module_name = self.source_module_name
+        hint = (
+            f" To quantize it, set a compatible granularity for '{module_name}' "
+            f"via module_name_configs."
+            if module_name is not None
+            else ""
+        )
         logger.warning(
-            "Tensor (target: %s) incompatible with block size "
-            "configuration: %s. Skipping quantization.",
+            "Tensor '%s' (target: %s, shape: %s) incompatible with block size "
+            "configuration: %s. Skipping quantization.%s",
+            self.source_name,
             self.quantization_target,
+            tuple(shape),
             error,
+            hint,
         )
         self._disabled.fill_(True)
 
@@ -156,7 +172,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
                 try:
                     scale, zero_point, minval = self.qparams_calculator(tensor)
                 except _BlockSizeMismatchError as e:
-                    self._warn_and_disable(e)
+                    self._warn_and_disable(e, tensor.shape)
                     return tensor
         else:
             # When the observer is not enabled, call the get_qparams
diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py
index 58ff74c5..4f2f925c 100644
--- a/tests/palettization/test_kmeans_palettizer.py
+++ b/tests/palettization/test_kmeans_palettizer.py
@@ -300,6 +300,38 @@ def test_disabled_fake_palett_removed_after_prepare(self, caplog):
         assert isinstance(prepared_model[2].parametrizations.weight[0], _FakePalettizeImplBase)
         assert not prepared_model[2].parametrizations.weight[0].is_disabled()
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The skip warning identifies the weight by FQN and shape."""
+        # axis 0 is out_features: 12 % 8 != 0 for the first Linear, 16 % 8 == 0
+        # for the second, so exactly one layer is skipped.
+        model = nn.Sequential(
+            nn.Linear(8, 12),
+            nn.ReLU(),
+            nn.Linear(12, 16),
+        )
+        example_inputs = (torch.randn(1, 8),)
+
+        config = KMeansPalettizerConfig(
+            global_config=ModuleKMeansPalettizerConfig(
+                op_state_spec={
+                    "weight": PalettizationSpec(
+                        n_bits=2,
+                        granularity=PerGroupedChannelGranularity(axis=0, group_size=8),
+                    )
+                },
+            )
+        )
+
+        with caplog.at_level(logging.WARNING):
+            KMeansPalettizer(model, config).prepare(example_inputs)
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping palettization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        message = skip_messages[0]
+
+        assert "'0.weight'" in message, message
+        assert "(12, 8)" in message, message
+
     def test_prepared_model_supports_torch_inference(
         self, simple_conv_linear_model, simple_model_input
     ):
diff --git a/tests/quantization/test_eager_quant.py b/tests/quantization/test_eager_quant.py
index 17049bbe..558fa3d0 100644
--- a/tests/quantization/test_eager_quant.py
+++ b/tests/quantization/test_eager_quant.py
@@ -3479,6 +3479,43 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog):
         test_input = torch.randn(1, 768)
         torch.testing.assert_close(prepared_model(test_input), ref_model(test_input))
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The skip warning identifies the weight by FQN and shape.
+
+        Without this the user has to hunt for the offending layer themselves.
+        """
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = nn.Sequential(nn.Linear(8, 12))
+        example_inputs = (torch.randn(1, 8),)
+
+        config = QuantizerConfig(
+            global_config=ModuleQuantizerConfig(
+                op_state_spec={
+                    "weight": QuantizationSpec(
+                        dtype="int8",
+                        qscheme="symmetric",
+                        granularity=PerBlockGranularity(axis=0, block_size=8),
+                    )
+                },
+                op_input_spec=None,
+                op_output_spec=None,
+            ),
+            execution_mode="eager",
+        )
+
+        with caplog.at_level(logging.WARNING):
+            Quantizer(model, config).prepare(example_inputs)
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        message = skip_messages[0]
+
+        # "0" is the Sequential index, so the weight FQN is "0.weight".
+        assert "'0.weight'" in message, message
+        assert "(12, 8)" in message, message
+        # The module name is the key the user needs for module_name_configs.
+        assert "'0' via module_name_configs" in message, message
+
     def test_divisible_block_size_not_disabled(self):
         """Linear(768, 1024) with block_size=32 on axis=0: 1024 % 32 == 0."""
         model = nn.Sequential(nn.Linear(768, 1024))
diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py
index 901cb24d..649e93e8 100644
--- a/tests/quantization/test_graph_mode_quantizer.py
+++ b/tests/quantization/test_graph_mode_quantizer.py
@@ -1490,6 +1490,42 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog, target, expected
         # The graph must still be runnable after the disabled nodes were removed.
         assert prepared_model(*example_inputs).shape == (1, 1024)
 
+    def test_skip_warning_names_the_offending_weight(self, caplog):
+        """The weight skip warning identifies the weight by FQN and shape.
+
+        In graph mode the name comes from the ``get_attr`` node feeding the
+        fake-quantize node, whose target is the dotted parameter FQN.
+        """
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = torch.nn.Sequential(torch.nn.Linear(8, 12))
+        example_inputs = (torch.randn(1, 8),)
+
+        config = QuantizerConfig(
+            global_config=ModuleQuantizerConfig(
+                op_state_spec={
+                    "weight": QuantizationSpec(
+                        dtype="int8",
+                        qscheme="symmetric",
+                        granularity=PerBlockGranularity(axis=0, block_size=8),
+                    )
+                },
+                op_input_spec=None,
+                op_output_spec=None,
+            ),
+            execution_mode="graph",
+        )
+
+        with caplog.at_level(logging.WARNING):
+            Quantizer(model, config).prepare(example_inputs)
+
+        skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
+        assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
+        message = skip_messages[0]
+
+        assert "'0.weight'" in message, message
+        assert "(12, 8)" in message, message
+        assert "'0' via module_name_configs" in message, message
+
     @pytest.mark.parametrize(
         "backend",
         [

From 2ebc604d1f3205ba1928f9d93fad5f768de26a17 Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Fri, 7 Aug 2026 11:54:13 -0700
Subject: [PATCH 02/10] cleanup

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 .../config/spec/compression_simulator.py      | 33 +++++---------
 src/coreai_opt/palettization/_source_names.py | 20 +++------
 .../palettization/kmeans/palettizer.py        |  5 +--
 .../palettization/spec/fake_palettize.py      | 11 +----
 .../quantization/_eager/quantizer.py          |  3 +-
 .../quantization/_graph/quantizer.py          |  3 +-
 src/coreai_opt/quantization/_source_names.py  | 44 +++++++------------
 .../quantization/spec/fake_quantize.py        |  8 ++--
 8 files changed, 43 insertions(+), 84 deletions(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index 2f4ac427..cd112965 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -26,27 +26,22 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     compression simulation is performed during training.
     """
 
-    # Model-level identity of the tensor this simulator compresses, recorded
-    # during prepare() so that forward-time diagnostics can name the tensor.
-    # These are class-level defaults rather than __init__ assignments so that
-    # subclass constructor chains stay untouched. Being plain strings, they are
-    # kept out of state_dict by nn.Module.__setattr__.
+    # Recorded during prepare() so that forward-time diagnostics can name the
+    # tensor. Class-level defaults keep this out of __init__, and out of
+    # state_dict, without touching subclass constructor chains.
     _source_module_name: str | None = None
     _source_param_name: str | None = None
 
     def set_source_name(self, module_name: str, param_name: str | None = None) -> None:
         """Record which model tensor this simulator compresses.
 
-        A simulator cannot discover its own position in the model from inside
-        ``forward``, so callers record it during ``prepare()`` instead. When one
-        simulator is shared by several modules, the first name recorded wins, so
-        that the reported name is stable across runs.
+        The first name recorded wins, so simulators shared by several modules
+        report a stable name.
 
         Args:
-            module_name: Fully-qualified name of the owning module as it appears
-                in ``named_modules()``, or ``""`` for a root-module parameter.
-            param_name: Local parameter name (e.g. ``"weight"``). ``None`` when
-                the simulator does not act on a named parameter.
+            module_name: Owning module's name in ``named_modules()``, or ``""``
+                for a root-module parameter.
+            param_name: Local parameter name (e.g. ``"weight"``), or ``None``.
         """
         if self._source_module_name is not None or self._source_param_name is not None:
             return
@@ -55,20 +50,16 @@ def set_source_name(self, module_name: str, param_name: str | None = None) -> No
 
     @property
     def source_name(self) -> str:
-        """FQN of the compressed tensor (e.g. ``"layers.0.q_proj.weight"``).
-
-        Returns ``"<unknown>"`` when no name was recorded, which is the case for
-        activations and for any simulator created outside a ``prepare()`` pass.
-        """
+        """FQN of the compressed tensor, or ``"<unknown>"`` if never recorded."""
         parts = [part for part in (self._source_module_name, self._source_param_name) if part]
         return ".".join(parts) if parts else "<unknown>"
 
     @property
     def source_module_name(self) -> str | None:
-        """FQN of the owning module, suitable as a ``module_name_configs`` key.
+        """Owning module's FQN, usable as a ``module_name_configs`` key.
 
-        ``None`` when no name was recorded or when the tensor belongs to the
-        root module (which cannot be addressed by module name).
+        ``None`` if never recorded, or for a root-module tensor that no module
+        name addresses.
         """
         return self._source_module_name or None
 
diff --git a/src/coreai_opt/palettization/_source_names.py b/src/coreai_opt/palettization/_source_names.py
index 914033e1..6c0fd687 100644
--- a/src/coreai_opt/palettization/_source_names.py
+++ b/src/coreai_opt/palettization/_source_names.py
@@ -3,17 +3,11 @@
 # 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
 
-"""Record which model tensor each fake-palettize module palettizes.
+"""Record which parameter each fake-palettize module palettizes.
 
-``_FakePalettizeImplBase`` warns and disables itself from inside ``forward``
-when a tensor is incompatible with the configured granularity or cluster
-dimension, but a module cannot discover its own position in the model at that
-point. The pass here runs during ``prepare()``, before centroids are computed,
-and stamps each fake-palettize module with the FQN of the parameter it
-palettizes so the warning can name the offending weight.
-
-Palettization targets weights only and is eager-only, so a single walk over
-``named_modules()`` covers every case.
+``_FakePalettizeImplBase`` warns and disables itself from inside ``forward``, where a
+module cannot know its own position in the model. The pass here runs during
+``prepare()``, before centroids are computed, so the warning can name the weight.
 """
 
 from __future__ import annotations
@@ -27,16 +21,14 @@
 def record_weight_source_names(model: nn.Module) -> None:
     """Stamp fake-palettize modules with the FQN of the parameter they palettize.
 
-    Fake-palettize modules live in the ``ParametrizationList`` registered for the
-    parameter they compress, so the owning module name and the parameter name
-    both come straight from ``named_modules()``.
-
     Args:
         model (nn.Module): The prepared model.
     """
     for module_name, module in model.named_modules(remove_duplicate=True):
         if not P.is_parametrized(module):
             continue
+        # A fake palettize lives in the ParametrizationList of the parameter it
+        # compresses.
         for param_name, parametrizations in module.parametrizations.items():
             for fake_palett in parametrizations:
                 if isinstance(fake_palett, _FakePalettizeImplBase):
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index c66676c8..ae12c50e 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -198,9 +198,8 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
-        # Record each fake-palettize module's parameter FQN so that the
-        # incompatibility warnings raised while computing centroids below can
-        # name the offending weight.
+        # Name fake palettize modules before centroid calculation below can
+        # warn about them.
         _record_weight_source_names(prepared_model)
 
         # Save example inputs for later use in calibration
diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py
index c82ac31f..ea342c65 100644
--- a/src/coreai_opt/palettization/spec/fake_palettize.py
+++ b/src/coreai_opt/palettization/spec/fake_palettize.py
@@ -79,20 +79,13 @@ def is_disabled(self) -> bool:
         return self._disabled
 
     def _warn_and_disable(self, setting: str, error: Exception, shape: torch.Size) -> None:
-        """Log a warning naming the offending weight and disable this module.
-
-        Names the tensor and its shape so that the palettization config can be
-        corrected without having to hunt for the layer first. The name is
-        recorded during ``prepare()``; it falls back to ``"<unknown>"`` for
-        modules created outside that path.
-        """
+        """Log a warning naming the offending weight and disable this module."""
         logger.warning(
             "Tensor '%s' (shape: %s) incompatible with %s: %s. Skipping palettization.",
             self.source_name,
             tuple(shape),
             setting,
-            # The granularity and cluster_dim messages end in a period; drop it
-            # so the sentence this builds has exactly one.
+            # These errors end in a period; drop it to avoid doubling up.
             str(error).rstrip("."),
         )
         self._disabled = True
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 35d49458..63c6567b 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -226,8 +226,7 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
-        # Record each weight FQ's parameter FQN so that the block-size warning
-        # raised during the upcoming forward pass can name the offending weight.
+        # Name weight FQs before the forward pass below can warn about them.
         _record_weight_source_names(model)
 
         _apply_weight_axis_defaults(model)
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index b259505a..156b7293 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -1301,8 +1301,7 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
-        # Record each weight FQ's parameter FQN so that the block-size warning
-        # raised during the upcoming forward pass can name the offending weight.
+        # Name weight FQs before the forward pass below can warn about them.
         _record_weight_source_names(model)
 
         # Apply weight axis defaults for per channel and per block quantization
diff --git a/src/coreai_opt/quantization/_source_names.py b/src/coreai_opt/quantization/_source_names.py
index 1e561959..feafe70e 100644
--- a/src/coreai_opt/quantization/_source_names.py
+++ b/src/coreai_opt/quantization/_source_names.py
@@ -3,22 +3,16 @@
 # 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
 
-"""Record which model tensor each weight fake-quantize module quantizes.
+"""Record which parameter each weight fake-quantize module quantizes.
 
-``FakeQuantizeImplBase`` warns and disables itself from inside ``forward`` when a
-tensor is incompatible with the configured block size, but a module cannot
-discover its own position in the model at that point. The passes here run during
-``prepare()``, before the first forward pass, and stamp each weight
-fake-quantize module with the FQN of the parameter it quantizes so the warning
-can name the offending weight.
+``FakeQuantizeImplBase`` warns and disables itself from inside ``forward``, where a
+module cannot know its own position in the model. These passes run during
+``prepare()``, before the first forward, so the warning can name the weight.
 
-Weights only: activation fake-quantize modules have no backing parameter, so
-they keep the unnamed warning.
-
-Like :mod:`coreai_opt.quantization._axis_defaults`, this module deliberately
-implements its own graph and eager walks instead of reusing the helpers in
-``_graph``/``_eager``. Both mode-specific quantizers import it, so importing
-either subpackage from here would create a circular import.
+Like :mod:`coreai_opt.quantization._axis_defaults`, this module implements its own
+graph and eager walks rather than reusing the ``_graph``/``_eager`` helpers: both
+mode-specific quantizers import it, so importing either subpackage here would be a
+circular import.
 """
 
 from __future__ import annotations
@@ -34,11 +28,6 @@
 def record_weight_source_names_graph(model: GraphModule) -> None:
     """Stamp weight fake-quantize modules in a graph-mode ``GraphModule``.
 
-    A weight fake-quantize node takes its value from the ``get_attr`` node for
-    the parameter, whose target is the dotted parameter FQN (e.g.
-    ``"layer1.0.weight"``), so the name needs no inference. This mirrors
-    ``_graph._prepare_for_export._get_weight_input_names``.
-
     Args:
         model (GraphModule): The prepared graph-mode ``GraphModule``.
     """
@@ -53,31 +42,30 @@ def record_weight_source_names_graph(model: GraphModule) -> None:
         if fake_quant.quantization_target != CompressionTargetTensor.WEIGHT:
             continue
 
-        # Activation fake-quantize nodes read from an op rather than a parameter.
+        # An already-compressed weight reaches the fake quantize through a
+        # decompression op (e.g. coreai.lut_to_dense) instead of a get_attr, and
+        # then carries no parameter name. See is_coreai_compressed_state_node.
         input_node = node.args[0]
         if input_node.op != "get_attr":
             continue
 
-        # "layer1.0.weight" -> ("layer1.0", "weight"); a root-module parameter
-        # such as "weight" has no module part.
-        target_path = str(input_node.target)
-        module_name, _, param_name = target_path.rpartition(".")
+        # A get_attr target is the dotted parameter FQN: "layer1.0.weight" ->
+        # ("layer1.0", "weight"). A root-module parameter has no module part.
+        module_name, _, param_name = str(input_node.target).rpartition(".")
         fake_quant.set_source_name(module_name, param_name)
 
 
 def record_weight_source_names_eager(model: nn.Module) -> None:
     """Stamp weight fake-quantize modules in an eager-mode model.
 
-    Weight fake-quantize modules live in the ``ParametrizationList`` registered
-    for the parameter they quantize, so the owning module name and the parameter
-    name both come straight from ``named_modules()``.
-
     Args:
         model (nn.Module): The prepared eager-mode model.
     """
     for module_name, module in model.named_modules(remove_duplicate=True):
         if not P.is_parametrized(module):
             continue
+        # A weight fake quantize lives in the ParametrizationList of the
+        # parameter it quantizes.
         for param_name, parametrizations in module.parametrizations.items():
             for fake_quant in parametrizations:
                 if not isinstance(fake_quant, FakeQuantizeImplBase):
diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py
index 8604176b..28eb0dfd 100644
--- a/src/coreai_opt/quantization/spec/fake_quantize.py
+++ b/src/coreai_opt/quantization/spec/fake_quantize.py
@@ -129,12 +129,10 @@ def enable_observer(self, enabled: bool = True) -> None:
         super().enable_observer(enabled)
 
     def _warn_and_disable(self, error: _BlockSizeMismatchError, shape: torch.Size) -> None:
-        """Log a warning and permanently disable this module.
+        """Log a warning naming the offending tensor and permanently disable this module.
 
-        Names the offending tensor and its shape so that the granularity can be
-        corrected without having to hunt for the layer first. The name is only
-        available for weights (recorded during ``prepare()``); activations fall
-        back to ``"<unknown>"``.
+        ``source_name`` is only recorded for weights, so activations report
+        ``"<unknown>"`` and get no config hint.
         """
         module_name = self.source_module_name
         hint = (

From a5227bd36bc7af7c85a18ac1d1566abc6072a75f Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Thu, 20 Aug 2026 09:56:19 -0700
Subject: [PATCH 03/10] refactor

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 .../config/spec/compression_simulator.py      | 67 ++++++++---------
 src/coreai_opt/palettization/_source_names.py | 35 ---------
 .../palettization/kmeans/palettizer.py        |  8 +-
 .../palettization/spec/fake_palettize.py      |  1 -
 .../quantization/_eager/quantizer.py          |  7 +-
 .../quantization/_graph/quantizer.py          |  7 +-
 src/coreai_opt/quantization/_source_names.py  | 75 -------------------
 .../quantization/spec/fake_quantize.py        | 16 +---
 tests/fixtures/quantization.py                | 11 ++-
 tests/palettization/test_kmeans_palettizer.py | 33 +++-----
 tests/quantization/test_eager_quant.py        | 38 +++-------
 .../quantization/test_graph_mode_quantizer.py | 37 +++------
 12 files changed, 83 insertions(+), 252 deletions(-)
 delete mode 100644 src/coreai_opt/palettization/_source_names.py
 delete mode 100644 src/coreai_opt/quantization/_source_names.py

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index cd112965..d138acef 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -9,6 +9,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
 
@@ -26,42 +27,9 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     compression simulation is performed during training.
     """
 
-    # Recorded during prepare() so that forward-time diagnostics can name the
-    # tensor. Class-level defaults keep this out of __init__, and out of
-    # state_dict, without touching subclass constructor chains.
-    _source_module_name: str | None = None
-    _source_param_name: str | None = None
-
-    def set_source_name(self, module_name: str, param_name: str | None = None) -> None:
-        """Record which model tensor this simulator compresses.
-
-        The first name recorded wins, so simulators shared by several modules
-        report a stable name.
-
-        Args:
-            module_name: Owning module's name in ``named_modules()``, or ``""``
-                for a root-module parameter.
-            param_name: Local parameter name (e.g. ``"weight"``), or ``None``.
-        """
-        if self._source_module_name is not None or self._source_param_name is not None:
-            return
-        self._source_module_name = module_name
-        self._source_param_name = param_name
-
-    @property
-    def source_name(self) -> str:
-        """FQN of the compressed tensor, or ``"<unknown>"`` if never recorded."""
-        parts = [part for part in (self._source_module_name, self._source_param_name) if part]
-        return ".".join(parts) if parts else "<unknown>"
-
-    @property
-    def source_module_name(self) -> str | None:
-        """Owning module's FQN, usable as a ``module_name_configs`` key.
-
-        ``None`` if never recorded, or for a root-module tensor that no module
-        name addresses.
-        """
-        return self._source_module_name or None
+    # FQN of the compressed tensor, recorded during prepare() because a module
+    # cannot discover its own name from forward().
+    source_name: str = "<unknown>"
 
     @abstractmethod
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
@@ -80,3 +48,30 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
             flowing through
         """
         pass
+
+
+def record_source_names_eager(model: nn.Module) -> None:
+    """Record the FQN of each parametrized tensor on the simulator compressing it."""
+    for module_name, module in model.named_modules():
+        if not P.is_parametrized(module):
+            continue
+        for param_name, parametrizations in module.parametrizations.items():
+            for simulator in parametrizations:
+                if isinstance(simulator, CompressionSimulatorBase):
+                    simulator.source_name = f"{module_name}.{param_name}"
+
+
+def record_source_names_graph(model: torch.fx.GraphModule) -> None:
+    """Record the FQN of each compressed parameter on the simulator compressing it.
+
+    A simulator reading from anything other than a ``get_attr`` node keeps the
+    default name: an activation, or a weight arriving through a decompression op,
+    has no parameter to name.
+    """
+    simulators = dict(model.named_modules(remove_duplicate=False))
+    for node in model.graph.nodes:
+        if node.op != "call_module" or not node.args:
+            continue
+        simulator = simulators.get(str(node.target))
+        if isinstance(simulator, CompressionSimulatorBase) and node.args[0].op == "get_attr":
+            simulator.source_name = str(node.args[0].target)
diff --git a/src/coreai_opt/palettization/_source_names.py b/src/coreai_opt/palettization/_source_names.py
deleted file mode 100644
index 6c0fd687..00000000
--- a/src/coreai_opt/palettization/_source_names.py
+++ /dev/null
@@ -1,35 +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
-
-"""Record which parameter each fake-palettize module palettizes.
-
-``_FakePalettizeImplBase`` warns and disables itself from inside ``forward``, where a
-module cannot know its own position in the model. The pass here runs during
-``prepare()``, before centroids are computed, so the warning can name the weight.
-"""
-
-from __future__ import annotations
-
-import torch.nn as nn
-import torch.nn.utils.parametrize as P
-
-from coreai_opt.palettization.spec.fake_palettize import _FakePalettizeImplBase
-
-
-def record_weight_source_names(model: nn.Module) -> None:
-    """Stamp fake-palettize modules with the FQN of the parameter they palettize.
-
-    Args:
-        model (nn.Module): The prepared model.
-    """
-    for module_name, module in model.named_modules(remove_duplicate=True):
-        if not P.is_parametrized(module):
-            continue
-        # A fake palettize lives in the ParametrizationList of the parameter it
-        # compresses.
-        for param_name, parametrizations in module.parametrizations.items():
-            for fake_palett in parametrizations:
-                if isinstance(fake_palett, _FakePalettizeImplBase):
-                    fake_palett.set_source_name(module_name, param_name)
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index ae12c50e..1988a67b 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -34,9 +34,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.palettization._source_names import (
-    record_weight_source_names as _record_weight_source_names,
-)
+from coreai_opt.config.spec.compression_simulator import record_source_names_eager
 from coreai_opt.palettization.base_palettizer import _BasePalettizer
 from coreai_opt.palettization.config.palettization_config import (
     KMeansPalettizerConfig,
@@ -198,9 +196,7 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
-        # Name fake palettize modules before centroid calculation below can
-        # warn about them.
-        _record_weight_source_names(prepared_model)
+        record_source_names_eager(prepared_model)
 
         # Save example inputs for later use in calibration
         self._example_inputs = tuple([ip.detach().clone() for ip in example_inputs])
diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py
index ea342c65..116c779c 100644
--- a/src/coreai_opt/palettization/spec/fake_palettize.py
+++ b/src/coreai_opt/palettization/spec/fake_palettize.py
@@ -85,7 +85,6 @@ def _warn_and_disable(self, setting: str, error: Exception, shape: torch.Size) -
             self.source_name,
             tuple(shape),
             setting,
-            # These errors end in a period; drop it to avoid doubling up.
             str(error).rstrip("."),
         )
         self._disabled = True
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 63c6567b..6d5fca1c 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -31,6 +31,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
+from coreai_opt.config.spec.compression_simulator import record_source_names_eager
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -39,9 +40,6 @@
     disable_activation_fake_quant,
     enable_weight_fake_quant,
 )
-from coreai_opt.quantization._source_names import (
-    record_weight_source_names_eager as _record_weight_source_names,
-)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     ModuleQuantizerConfig,
@@ -226,8 +224,7 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
-        # Name weight FQs before the forward pass below can warn about them.
-        _record_weight_source_names(model)
+        record_source_names_eager(model)
 
         _apply_weight_axis_defaults(model)
         validate_activation_axes(model)
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 156b7293..7ebd3a7c 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -46,6 +46,7 @@
 from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
+from coreai_opt.config.spec.compression_simulator import record_source_names_graph
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -54,9 +55,6 @@
     disable_activation_fake_quant,
     enable_weight_fake_quant,
 )
-from coreai_opt.quantization._source_names import (
-    record_weight_source_names_graph as _record_weight_source_names,
-)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     KVCacheQuantConfig,
@@ -1301,8 +1299,7 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
-        # Name weight FQs before the forward pass below can warn about them.
-        _record_weight_source_names(model)
+        record_source_names_graph(model)
 
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)
diff --git a/src/coreai_opt/quantization/_source_names.py b/src/coreai_opt/quantization/_source_names.py
deleted file mode 100644
index feafe70e..00000000
--- a/src/coreai_opt/quantization/_source_names.py
+++ /dev/null
@@ -1,75 +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
-
-"""Record which parameter each weight fake-quantize module quantizes.
-
-``FakeQuantizeImplBase`` warns and disables itself from inside ``forward``, where a
-module cannot know its own position in the model. These passes run during
-``prepare()``, before the first forward, so the warning can name the weight.
-
-Like :mod:`coreai_opt.quantization._axis_defaults`, this module implements its own
-graph and eager walks rather than reusing the ``_graph``/``_eager`` helpers: both
-mode-specific quantizers import it, so importing either subpackage here would be a
-circular import.
-"""
-
-from __future__ import annotations
-
-import torch.nn as nn
-import torch.nn.utils.parametrize as P
-from torch.fx import GraphModule
-
-from coreai_opt.config.spec import CompressionTargetTensor
-from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
-
-
-def record_weight_source_names_graph(model: GraphModule) -> None:
-    """Stamp weight fake-quantize modules in a graph-mode ``GraphModule``.
-
-    Args:
-        model (GraphModule): The prepared graph-mode ``GraphModule``.
-    """
-    modules = dict(model.named_modules(remove_duplicate=False))
-
-    for node in model.graph.nodes:
-        if node.op != "call_module":
-            continue
-        fake_quant = modules.get(str(node.target))
-        if not isinstance(fake_quant, FakeQuantizeImplBase):
-            continue
-        if fake_quant.quantization_target != CompressionTargetTensor.WEIGHT:
-            continue
-
-        # An already-compressed weight reaches the fake quantize through a
-        # decompression op (e.g. coreai.lut_to_dense) instead of a get_attr, and
-        # then carries no parameter name. See is_coreai_compressed_state_node.
-        input_node = node.args[0]
-        if input_node.op != "get_attr":
-            continue
-
-        # A get_attr target is the dotted parameter FQN: "layer1.0.weight" ->
-        # ("layer1.0", "weight"). A root-module parameter has no module part.
-        module_name, _, param_name = str(input_node.target).rpartition(".")
-        fake_quant.set_source_name(module_name, param_name)
-
-
-def record_weight_source_names_eager(model: nn.Module) -> None:
-    """Stamp weight fake-quantize modules in an eager-mode model.
-
-    Args:
-        model (nn.Module): The prepared eager-mode model.
-    """
-    for module_name, module in model.named_modules(remove_duplicate=True):
-        if not P.is_parametrized(module):
-            continue
-        # A weight fake quantize lives in the ParametrizationList of the
-        # parameter it quantizes.
-        for param_name, parametrizations in module.parametrizations.items():
-            for fake_quant in parametrizations:
-                if not isinstance(fake_quant, FakeQuantizeImplBase):
-                    continue
-                if fake_quant.quantization_target != CompressionTargetTensor.WEIGHT:
-                    continue
-                fake_quant.set_source_name(module_name, param_name)
diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py
index 28eb0dfd..f02db512 100644
--- a/src/coreai_opt/quantization/spec/fake_quantize.py
+++ b/src/coreai_opt/quantization/spec/fake_quantize.py
@@ -129,26 +129,14 @@ def enable_observer(self, enabled: bool = True) -> None:
         super().enable_observer(enabled)
 
     def _warn_and_disable(self, error: _BlockSizeMismatchError, shape: torch.Size) -> None:
-        """Log a warning naming the offending tensor and permanently disable this module.
-
-        ``source_name`` is only recorded for weights, so activations report
-        ``"<unknown>"`` and get no config hint.
-        """
-        module_name = self.source_module_name
-        hint = (
-            f" To quantize it, set a compatible granularity for '{module_name}' "
-            f"via module_name_configs."
-            if module_name is not None
-            else ""
-        )
+        """Log a warning naming the offending tensor and permanently disable this module."""
         logger.warning(
             "Tensor '%s' (target: %s, shape: %s) incompatible with block size "
-            "configuration: %s. Skipping quantization.%s",
+            "configuration: %s. Skipping quantization.",
             self.source_name,
             self.quantization_target,
             tuple(shape),
             error,
-            hint,
         )
         self._disabled.fill_(True)
 
diff --git a/tests/fixtures/quantization.py b/tests/fixtures/quantization.py
index 4ee165b0..0b801104 100644
--- a/tests/fixtures/quantization.py
+++ b/tests/fixtures/quantization.py
@@ -17,6 +17,7 @@
     PerBlockGranularity,
     PerChannelGranularity,
     PerTensorGranularity,
+    QuantizationGranularity,
     QuantizationScheme,
     QuantizationSpec,
 )
@@ -50,23 +51,27 @@ def make_quant_config(
     weight_dtype: torch.dtype | str | None,
     act_dtype: torch.dtype | str | None,
     execution_mode: str,
+    granularity: QuantizationGranularity | None = None,
 ) -> QuantizerConfig:
-    """Build a per-tensor symmetric QuantizerConfig for export tests.
+    """Build a symmetric QuantizerConfig for export tests.
 
     Args:
         weight_dtype (torch.dtype | str | None): Weight dtype, or None to disable.
         act_dtype (torch.dtype | str | None): Activation dtype, or None to disable.
         execution_mode (str): Either "eager" or "graph".
+        granularity (QuantizationGranularity | None): Granularity for both the
+            weight and activation specs. Defaults to per-tensor.
 
     Returns:
-        QuantizerConfig: Config with the requested per-tensor symmetric specs.
+        QuantizerConfig: Config with the requested symmetric specs.
     """
+    granularity = granularity or PerTensorGranularity()
 
     def _spec(dtype: torch.dtype | str) -> QuantizationSpec:
         return QuantizationSpec(
             dtype=dtype,
             qscheme=QuantizationScheme.SYMMETRIC,
-            granularity=PerTensorGranularity(),
+            granularity=granularity,
         )
 
     weight_spec = _spec(weight_dtype) if weight_dtype is not None else None
diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py
index 4f2f925c..bbad7bad 100644
--- a/tests/palettization/test_kmeans_palettizer.py
+++ b/tests/palettization/test_kmeans_palettizer.py
@@ -302,35 +302,26 @@ def test_disabled_fake_palett_removed_after_prepare(self, caplog):
 
     def test_skip_warning_names_the_offending_weight(self, caplog):
         """The skip warning identifies the weight by FQN and shape."""
-        # axis 0 is out_features: 12 % 8 != 0 for the first Linear, 16 % 8 == 0
-        # for the second, so exactly one layer is skipped.
-        model = nn.Sequential(
-            nn.Linear(8, 12),
-            nn.ReLU(),
-            nn.Linear(12, 16),
+        # axis 0 is out_features: 12 % 8 != 0.
+        model = nn.Sequential(nn.Linear(8, 12))
+        spec = PalettizationSpec(
+            n_bits=2, granularity=PerGroupedChannelGranularity(axis=0, group_size=8)
         )
-        example_inputs = (torch.randn(1, 8),)
-
         config = KMeansPalettizerConfig(
-            global_config=ModuleKMeansPalettizerConfig(
-                op_state_spec={
-                    "weight": PalettizationSpec(
-                        n_bits=2,
-                        granularity=PerGroupedChannelGranularity(axis=0, group_size=8),
-                    )
-                },
-            )
+            global_config=ModuleKMeansPalettizerConfig(op_state_spec={"weight": spec})
         )
 
         with caplog.at_level(logging.WARNING):
-            KMeansPalettizer(model, config).prepare(example_inputs)
+            KMeansPalettizer(model, config).prepare((torch.randn(1, 8),))
 
         skip_messages = [msg for msg in caplog.messages if "Skipping palettization" in msg]
         assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
-        message = skip_messages[0]
-
-        assert "'0.weight'" in message, message
-        assert "(12, 8)" in message, message
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (shape: (12, 8)) incompatible with granularity: Tensor size 12 "
+            "along axis 0 is not divisible by group_size 8. For per-grouped-channel "
+            "palettization, the tensor shape along the specified axis must be divisible by "
+            "group_size. Skipping palettization."
+        )
 
     def test_prepared_model_supports_torch_inference(
         self, simple_conv_linear_model, simple_model_input
diff --git a/tests/quantization/test_eager_quant.py b/tests/quantization/test_eager_quant.py
index 558fa3d0..b4de343a 100644
--- a/tests/quantization/test_eager_quant.py
+++ b/tests/quantization/test_eager_quant.py
@@ -38,6 +38,7 @@
     default_weight_quantization_spec,
 )
 from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+from tests.fixtures.quantization import make_quant_config
 
 
 class InnerModule(nn.Module):
@@ -3480,41 +3481,26 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog):
         torch.testing.assert_close(prepared_model(test_input), ref_model(test_input))
 
     def test_skip_warning_names_the_offending_weight(self, caplog):
-        """The skip warning identifies the weight by FQN and shape.
-
-        Without this the user has to hunt for the offending layer themselves.
-        """
+        """The skip warning identifies the weight by FQN and shape."""
         # axis 0 is out_features: 12 % 8 != 0.
         model = nn.Sequential(nn.Linear(8, 12))
-        example_inputs = (torch.randn(1, 8),)
-
-        config = QuantizerConfig(
-            global_config=ModuleQuantizerConfig(
-                op_state_spec={
-                    "weight": QuantizationSpec(
-                        dtype="int8",
-                        qscheme="symmetric",
-                        granularity=PerBlockGranularity(axis=0, block_size=8),
-                    )
-                },
-                op_input_spec=None,
-                op_output_spec=None,
-            ),
+        config = make_quant_config(
+            weight_dtype="int8",
+            act_dtype=None,
             execution_mode="eager",
+            granularity=PerBlockGranularity(axis=0, block_size=8),
         )
 
         with caplog.at_level(logging.WARNING):
-            Quantizer(model, config).prepare(example_inputs)
+            Quantizer(model, config).prepare((torch.randn(1, 8),))
 
         skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
         assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
-        message = skip_messages[0]
-
-        # "0" is the Sequential index, so the weight FQN is "0.weight".
-        assert "'0.weight'" in message, message
-        assert "(12, 8)" in message, message
-        # The module name is the key the user needs for module_name_configs.
-        assert "'0' via module_name_configs" in message, message
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (target: CompressionTargetTensor.WEIGHT, shape: (12, 8)) "
+            "incompatible with block size configuration: Tensor size 12 along axis 0 is not "
+            "divisible by block size 8. Skipping quantization."
+        )
 
     def test_divisible_block_size_not_disabled(self):
         """Linear(768, 1024) with block_size=32 on axis=0: 1024 % 32 == 0."""
diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py
index 649e93e8..4dfcbc88 100644
--- a/tests/quantization/test_graph_mode_quantizer.py
+++ b/tests/quantization/test_graph_mode_quantizer.py
@@ -49,6 +49,7 @@
     MovingAverageQParamsCalculator,
     StaticQParamsCalculator,
 )
+from tests.fixtures.quantization import make_quant_config
 from tests.models.simple import SimpleModel
 
 
@@ -1491,40 +1492,26 @@ def test_non_divisible_block_size_warns_and_skips(self, caplog, target, expected
         assert prepared_model(*example_inputs).shape == (1, 1024)
 
     def test_skip_warning_names_the_offending_weight(self, caplog):
-        """The weight skip warning identifies the weight by FQN and shape.
-
-        In graph mode the name comes from the ``get_attr`` node feeding the
-        fake-quantize node, whose target is the dotted parameter FQN.
-        """
+        """The skip warning identifies the weight by FQN and shape."""
         # axis 0 is out_features: 12 % 8 != 0.
         model = torch.nn.Sequential(torch.nn.Linear(8, 12))
-        example_inputs = (torch.randn(1, 8),)
-
-        config = QuantizerConfig(
-            global_config=ModuleQuantizerConfig(
-                op_state_spec={
-                    "weight": QuantizationSpec(
-                        dtype="int8",
-                        qscheme="symmetric",
-                        granularity=PerBlockGranularity(axis=0, block_size=8),
-                    )
-                },
-                op_input_spec=None,
-                op_output_spec=None,
-            ),
+        config = make_quant_config(
+            weight_dtype="int8",
+            act_dtype=None,
             execution_mode="graph",
+            granularity=PerBlockGranularity(axis=0, block_size=8),
         )
 
         with caplog.at_level(logging.WARNING):
-            Quantizer(model, config).prepare(example_inputs)
+            Quantizer(model, config).prepare((torch.randn(1, 8),))
 
         skip_messages = [msg for msg in caplog.messages if "Skipping quantization" in msg]
         assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
-        message = skip_messages[0]
-
-        assert "'0.weight'" in message, message
-        assert "(12, 8)" in message, message
-        assert "'0' via module_name_configs" in message, message
+        assert skip_messages[0] == (
+            "Tensor '0.weight' (target: CompressionTargetTensor.WEIGHT, shape: (12, 8)) "
+            "incompatible with block size configuration: Tensor size 12 along axis 0 is not "
+            "divisible by block size 8. Skipping quantization."
+        )
 
     @pytest.mark.parametrize(
         "backend",

From 59b4b07d7e3cff99ccf1d27029904f321df2d41b Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 13:22:57 -0700
Subject: [PATCH 04/10] nit

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 src/coreai_opt/config/spec/compression_simulator.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index d138acef..8d79ca08 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -27,8 +27,7 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     compression simulation is performed during training.
     """
 
-    # FQN of the compressed tensor, recorded during prepare() because a module
-    # cannot discover its own name from forward().
+    # FQN of the compressed tensor
     source_name: str = "<unknown>"
 
     @abstractmethod

From 007c8641b2d28901cf79b8c22f8d6a17076b8cdc Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 14:15:19 -0700
Subject: [PATCH 05/10] refactor

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 .../config/spec/compression_simulator.py      |  3 +-
 .../kmeans/kmeans_fake_palettize.py           |  6 +++-
 .../palettization/kmeans/palettizer.py        | 34 ++++++-------------
 .../palettization/spec/fake_palettize.py      | 11 ------
 tests/palettization/test_kmeans_palettizer.py |  4 +--
 5 files changed, 20 insertions(+), 38 deletions(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index 8d79ca08..77012ada 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -55,9 +55,10 @@ def record_source_names_eager(model: nn.Module) -> None:
         if not P.is_parametrized(module):
             continue
         for param_name, parametrizations in module.parametrizations.items():
+            name = f"{module_name}.{param_name}" if module_name else param_name
             for simulator in parametrizations:
                 if isinstance(simulator, CompressionSimulatorBase):
-                    simulator.source_name = f"{module_name}.{param_name}"
+                    simulator.source_name = name
 
 
 def record_source_names_graph(model: torch.fx.GraphModule) -> None:
diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
index 566c140a..ef69f920 100644
--- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
+++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
@@ -166,7 +166,11 @@ def ensure_initialized(self, tensor: torch.Tensor) -> None:
             self._initialize(tensor.detach())
         except (_IncompatibleClusterDimError, _IncompatibleGranularityError) as e:
             logger.warning(
-                f"Tensor incompatible with configured spec: {e}. Skipping palettization."
+                "Tensor '%s' (shape: %s) incompatible with configured spec: %s. "
+                "Skipping palettization.",
+                self.source_name,
+                tuple(tensor.shape),
+                str(e).rstrip("."),
             )
             self._disabled = True
 
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index f0947f33..ef058f68 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -65,19 +65,14 @@ class _FakePalettInfo:
     """Metadata about _KMeansFakePalettize module and its associated module / param"""
 
     module: torch.nn.Module
-    module_name: str
     attr_name: str
     idx: int
     fp_module: _KMeansFakePalettize
     weight: torch.Tensor
 
-    @property
-    def layer_name(self) -> str:
-        return f"{self.module_name}.{self.attr_name}" if self.module_name else self.attr_name
-
 
 def _calculate_centroids_for_module(
-    args: tuple[_KMeansFakePalettize, torch.Tensor, str],
+    args: tuple[_KMeansFakePalettize, torch.Tensor],
 ) -> _KMeansFakePalettize:
     """Compute centroids for a single _KMeansFakePalettize module.
 
@@ -88,20 +83,22 @@ def _calculate_centroids_for_module(
     Invokes ``fp_module.forward(weight)`` to mirror the sequential path.
 
     Args:
-        args (tuple[_KMeansFakePalettize, torch.Tensor, str]): Tuple of
-            ``(fp_module, weight, layer_name)``.
+        args (tuple[_KMeansFakePalettize, torch.Tensor]): Tuple of
+            ``(fp_module, weight)``.
 
     Returns:
         _KMeansFakePalettize: The mutated module, ready to be swapped into
             the parent's parametrization slot.
     """
-    fp_module, weight, layer_name = args
+    fp_module, weight = args
 
     try:
         with torch.no_grad():
             fp_module(weight)
     except Exception as e:
-        raise RuntimeError(f"Centroid calculation failed for layer {layer_name!r}") from e
+        raise RuntimeError(
+            f"Centroid calculation failed for weight {fp_module.source_name!r}"
+        ) from e
 
     return fp_module
 
@@ -508,7 +505,7 @@ def _collect_fake_palett_info(self, *, to_cpu: bool) -> list[_FakePalettInfo]:
         shipped to a spawned worker process.
         """
         fp_info: list[_FakePalettInfo] = []
-        for module_name, module in self._model.named_modules(remove_duplicate=True):
+        for _module_name, module in self._model.named_modules(remove_duplicate=True):
             if not P.is_parametrized(module):
                 continue
             for attr_name, parametrizations in module.parametrizations.items():
@@ -520,7 +517,6 @@ def _collect_fake_palett_info(self, *, to_cpu: bool) -> list[_FakePalettInfo]:
                         fp_info.append(
                             _FakePalettInfo(
                                 module=module,
-                                module_name=module_name,
                                 attr_name=attr_name,
                                 idx=idx,
                                 fp_module=p,
@@ -544,7 +540,7 @@ def _calculate_centroids_sequential(self) -> None:
             return
 
         results = [
-            _calculate_centroids_for_module((info.fp_module, info.weight, info.layer_name))
+            _calculate_centroids_for_module((info.fp_module, info.weight))
             for info in tqdm(fp_info, desc="Palettizing layers (num_workers=1)")
         ]
         self._apply_centroid_results(fp_info, results)
@@ -564,7 +560,7 @@ def _calculate_centroids_parallel(self, num_workers: int) -> None:
 
         # spawn (not fork) so workers don't inherit the parent's CUDA context
         # or other process-global state.
-        pool_args = [(info.fp_module, info.weight, info.layer_name) for info in fp_info]
+        pool_args = [(info.fp_module, info.weight) for info in fp_info]
         ctx = mp.get_context("spawn")
         with ctx.Pool(processes=effective_workers) as pool:
             results = list(
@@ -591,16 +587,8 @@ def _apply_centroid_results(
         (mutated in place), so the assignment is a harmless no-op there.
         """
         for info, new_fp in zip(fp_info, results, strict=True):
-<<<<<<< HEAD
-            if getattr(new_fp, "_disabled", False):
-                logger.warning("Disabling palettization for weight '%s'", new_fp.source_name)
-            # ParametrizationList supports item assignment; this swaps the
-            # worker's mutated module into the live model without touching
-            # the surrounding parametrization registration.
-=======
             if new_fp.is_disabled():
-                logger.warning("Disabling palettization for layer %r", info.layer_name)
->>>>>>> main
+                logger.warning("Disabling palettization for weight '%s'", new_fp.source_name)
             info.module.parametrizations[info.attr_name][info.idx] = new_fp
 
     @staticmethod
diff --git a/src/coreai_opt/palettization/spec/fake_palettize.py b/src/coreai_opt/palettization/spec/fake_palettize.py
index 857ca5a5..34dff090 100644
--- a/src/coreai_opt/palettization/spec/fake_palettize.py
+++ b/src/coreai_opt/palettization/spec/fake_palettize.py
@@ -62,17 +62,6 @@ def is_disabled(self) -> bool:
         """Return True if fake palettization has been disabled."""
         return self._disabled
 
-    def _warn_and_disable(self, setting: str, error: Exception, shape: torch.Size) -> None:
-        """Log a warning naming the offending weight and disable this module."""
-        logger.warning(
-            "Tensor '%s' (shape: %s) incompatible with %s: %s. Skipping palettization.",
-            self.source_name,
-            tuple(shape),
-            setting,
-            str(error).rstrip("."),
-        )
-        self._disabled = True
-
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """Fake-palettize ``tensor`` through the enable/disable lifecycle.
 
diff --git a/tests/palettization/test_kmeans_palettizer.py b/tests/palettization/test_kmeans_palettizer.py
index 2dfba56a..a152998e 100644
--- a/tests/palettization/test_kmeans_palettizer.py
+++ b/tests/palettization/test_kmeans_palettizer.py
@@ -313,8 +313,8 @@ def test_skip_warning_names_the_offending_weight(self, caplog):
         skip_messages = [msg for msg in caplog.messages if "Skipping palettization" in msg]
         assert len(skip_messages) == 1, f"Expected one skip warning, got {skip_messages}"
         assert skip_messages[0] == (
-            "Tensor '0.weight' (shape: (12, 8)) incompatible with granularity: Tensor size 12 "
-            "along axis 0 is not divisible by group_size 8. For per-grouped-channel "
+            "Tensor '0.weight' (shape: (12, 8)) incompatible with configured spec: Tensor size "
+            "12 along axis 0 is not divisible by group_size 8. For per-grouped-channel "
             "palettization, the tensor shape along the specified axis must be divisible by "
             "group_size. Skipping palettization."
         )

From 0535976bd97951ad5ba035bde2a361fb45921f5c Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 14:19:28 -0700
Subject: [PATCH 06/10] rename

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 src/coreai_opt/config/spec/compression_simulator.py    | 10 +++++-----
 .../palettization/kmeans/kmeans_fake_palettize.py      |  2 +-
 src/coreai_opt/palettization/kmeans/palettizer.py      |  8 ++++----
 src/coreai_opt/quantization/_eager/quantizer.py        |  4 ++--
 src/coreai_opt/quantization/_graph/quantizer.py        |  4 ++--
 src/coreai_opt/quantization/spec/fake_quantize.py      |  2 +-
 6 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index 77012ada..e0f0b623 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -28,7 +28,7 @@ class CompressionSimulatorBase(_ClassRegistryMixin, nn.Module):
     """
 
     # FQN of the compressed tensor
-    source_name: str = "<unknown>"
+    tensor_fqn: str = "<unknown>"
 
     @abstractmethod
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
@@ -49,7 +49,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         pass
 
 
-def record_source_names_eager(model: nn.Module) -> None:
+def record_tensor_fqns_eager(model: nn.Module) -> None:
     """Record the FQN of each parametrized tensor on the simulator compressing it."""
     for module_name, module in model.named_modules():
         if not P.is_parametrized(module):
@@ -58,10 +58,10 @@ def record_source_names_eager(model: nn.Module) -> None:
             name = f"{module_name}.{param_name}" if module_name else param_name
             for simulator in parametrizations:
                 if isinstance(simulator, CompressionSimulatorBase):
-                    simulator.source_name = name
+                    simulator.tensor_fqn = name
 
 
-def record_source_names_graph(model: torch.fx.GraphModule) -> None:
+def record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
     """Record the FQN of each compressed parameter on the simulator compressing it.
 
     A simulator reading from anything other than a ``get_attr`` node keeps the
@@ -74,4 +74,4 @@ def record_source_names_graph(model: torch.fx.GraphModule) -> None:
             continue
         simulator = simulators.get(str(node.target))
         if isinstance(simulator, CompressionSimulatorBase) and node.args[0].op == "get_attr":
-            simulator.source_name = str(node.args[0].target)
+            simulator.tensor_fqn = str(node.args[0].target)
diff --git a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
index ef69f920..18404b11 100644
--- a/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
+++ b/src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py
@@ -168,7 +168,7 @@ def ensure_initialized(self, tensor: torch.Tensor) -> None:
             logger.warning(
                 "Tensor '%s' (shape: %s) incompatible with configured spec: %s. "
                 "Skipping palettization.",
-                self.source_name,
+                self.tensor_fqn,
                 tuple(tensor.shape),
                 str(e).rstrip("."),
             )
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index ef058f68..2827c1c4 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -36,7 +36,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig, ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import record_source_names_eager
+from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_eager
 from coreai_opt.palettization.base_palettizer import _BasePalettizer
 from coreai_opt.palettization.config.palettization_config import (
     KMeansPalettizerConfig,
@@ -97,7 +97,7 @@ def _calculate_centroids_for_module(
             fp_module(weight)
     except Exception as e:
         raise RuntimeError(
-            f"Centroid calculation failed for weight {fp_module.source_name!r}"
+            f"Centroid calculation failed for weight {fp_module.tensor_fqn!r}"
         ) from e
 
     return fp_module
@@ -199,7 +199,7 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
-        record_source_names_eager(prepared_model)
+        record_tensor_fqns_eager(prepared_model)
 
         # Load precomputed sensitivities if provided
         if sensitivity_path is not None:
@@ -588,7 +588,7 @@ def _apply_centroid_results(
         """
         for info, new_fp in zip(fp_info, results, strict=True):
             if new_fp.is_disabled():
-                logger.warning("Disabling palettization for weight '%s'", new_fp.source_name)
+                logger.warning("Disabling palettization for weight '%s'", new_fp.tensor_fqn)
             info.module.parametrizations[info.attr_name][info.idx] = new_fp
 
     @staticmethod
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 6d5fca1c..98cfcc84 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -31,7 +31,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import record_source_names_eager
+from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_eager
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -224,7 +224,7 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
-        record_source_names_eager(model)
+        record_tensor_fqns_eager(model)
 
         _apply_weight_axis_defaults(model)
         validate_activation_axes(model)
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index ec1424e0..20d5e2c3 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -46,7 +46,7 @@
 from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
-from coreai_opt.config.spec.compression_simulator import record_source_names_graph
+from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_graph
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -1380,7 +1380,7 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
-        record_source_names_graph(model)
+        record_tensor_fqns_graph(model)
 
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)
diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py
index f02db512..31b4423e 100644
--- a/src/coreai_opt/quantization/spec/fake_quantize.py
+++ b/src/coreai_opt/quantization/spec/fake_quantize.py
@@ -133,7 +133,7 @@ def _warn_and_disable(self, error: _BlockSizeMismatchError, shape: torch.Size) -
         logger.warning(
             "Tensor '%s' (target: %s, shape: %s) incompatible with block size "
             "configuration: %s. Skipping quantization.",
-            self.source_name,
+            self.tensor_fqn,
             self.quantization_target,
             tuple(shape),
             error,

From b124fdf450c6a9d82b93027322c5af74159f3ea1 Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 14:29:12 -0700
Subject: [PATCH 07/10] add changelog

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 changelog.d/85.changed | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/85.changed

diff --git a/changelog.d/85.changed b/changelog.d/85.changed
new file mode 100644
index 00000000..11db1aea
--- /dev/null
+++ b/changelog.d/85.changed
@@ -0,0 +1 @@
+When a weight is incompatible with the configured quantization block size or palettization granularity, the warning now reports the weight's fully qualified name and shape along with the module name to use in `module_name_configs`, instead of only the compression target.

From 41c169b4e9120237f108bdac49bee89c3512ec34 Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 16:12:43 -0700
Subject: [PATCH 08/10] rename

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 src/coreai_opt/config/spec/compression_simulator.py | 4 ++--
 src/coreai_opt/palettization/kmeans/palettizer.py   | 4 ++--
 src/coreai_opt/quantization/_eager/quantizer.py     | 4 ++--
 src/coreai_opt/quantization/_graph/quantizer.py     | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index e0f0b623..bb6a1011 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -49,7 +49,7 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         pass
 
 
-def record_tensor_fqns_eager(model: nn.Module) -> None:
+def _record_tensor_fqns_eager(model: nn.Module) -> None:
     """Record the FQN of each parametrized tensor on the simulator compressing it."""
     for module_name, module in model.named_modules():
         if not P.is_parametrized(module):
@@ -61,7 +61,7 @@ def record_tensor_fqns_eager(model: nn.Module) -> None:
                     simulator.tensor_fqn = name
 
 
-def record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
+def _record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
     """Record the FQN of each compressed parameter on the simulator compressing it.
 
     A simulator reading from anything other than a ``get_attr`` node keeps the
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index 2827c1c4..99215dbc 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -36,7 +36,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig, ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_eager
+from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_eager
 from coreai_opt.palettization.base_palettizer import _BasePalettizer
 from coreai_opt.palettization.config.palettization_config import (
     KMeansPalettizerConfig,
@@ -199,7 +199,7 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
-        record_tensor_fqns_eager(prepared_model)
+        _record_tensor_fqns_eager(prepared_model)
 
         # Load precomputed sensitivities if provided
         if sensitivity_path is not None:
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 98cfcc84..9e76dfbe 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -31,7 +31,7 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_eager
+from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_eager
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -224,7 +224,7 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
-        record_tensor_fqns_eager(model)
+        _record_tensor_fqns_eager(model)
 
         _apply_weight_axis_defaults(model)
         validate_activation_axes(model)
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 20d5e2c3..8974962a 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -46,7 +46,7 @@
 from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
-from coreai_opt.config.spec.compression_simulator import record_tensor_fqns_graph
+from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_graph
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -1380,7 +1380,7 @@ def _postprocess_prepared_model(model: torch.fx.GraphModule) -> None:
         # granularity is mathematically incorrect. Force per-tensor.
         force_per_tensor_for_channel_altering_ops(model)
 
-        record_tensor_fqns_graph(model)
+        _record_tensor_fqns_graph(model)
 
         # Apply weight axis defaults for per channel and per block quantization
         _apply_weight_axis_defaults(model)

From 2142d0b35c7fee1c317a20aa22c83b57da98264f Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 16:22:47 -0700
Subject: [PATCH 09/10] refactor record utils to handler

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 .../insertion/torch_function/handler.py       | 17 +++++++++++
 .../config/spec/compression_simulator.py      | 28 -------------------
 .../palettization/kmeans/palettizer.py        |  3 --
 .../quantization/_eager/quantizer.py          |  3 --
 .../quantization/_graph/quantizer.py          | 18 +++++++++++-
 5 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/src/coreai_opt/_utils/insertion/torch_function/handler.py b/src/coreai_opt/_utils/insertion/torch_function/handler.py
index a5f1682c..9225a532 100644
--- a/src/coreai_opt/_utils/insertion/torch_function/handler.py
+++ b/src/coreai_opt/_utils/insertion/torch_function/handler.py
@@ -9,9 +9,11 @@
 from typing import Any
 
 import torch.nn as nn
+import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.torch_utils import NamedModule
 from coreai_opt.config import CompressionConfig
+from coreai_opt.config.spec import CompressionSimulatorBase
 
 from .base_supported_ops_registry import BaseSupportedOpsRegistry
 from .modes import (
@@ -21,6 +23,18 @@
 from .types import ModuleCompressionComponents
 
 
+def _record_tensor_fqns(model: nn.Module) -> None:
+    """Record the FQN of each parametrized tensor on the simulator compressing it."""
+    for module_name, module in model.named_modules():
+        if not P.is_parametrized(module):
+            continue
+        for param_name, parametrizations in module.parametrizations.items():
+            fqn = f"{module_name}.{param_name}" if module_name else param_name
+            for simulator in parametrizations:
+                if isinstance(simulator, CompressionSimulatorBase):
+                    simulator.tensor_fqn = fqn
+
+
 class TorchFunctionEagerHandler:
     """
     Prepares the model for compression by inserting weight and activation
@@ -70,6 +84,9 @@ def prepare(self, model: nn.Module, example_inputs: tuple[Any, ...]) -> nn.Modul
         # seen again later in the forward pass.
         register_optimization_mode.register_all_activations()
         register_optimization_mode.register_all_states()
+
+        _record_tensor_fqns(model)
+
         if self._is_weight_only_optimization(self.module_components_dict):
             return model
 
diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index bb6a1011..980f8b6c 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -9,7 +9,6 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.utils.parametrize as P
 
 from coreai_opt._utils.registry_utils import ClassRegistryMixin as _ClassRegistryMixin
 
@@ -48,30 +47,3 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """
         pass
 
-
-def _record_tensor_fqns_eager(model: nn.Module) -> None:
-    """Record the FQN of each parametrized tensor on the simulator compressing it."""
-    for module_name, module in model.named_modules():
-        if not P.is_parametrized(module):
-            continue
-        for param_name, parametrizations in module.parametrizations.items():
-            name = f"{module_name}.{param_name}" if module_name else param_name
-            for simulator in parametrizations:
-                if isinstance(simulator, CompressionSimulatorBase):
-                    simulator.tensor_fqn = name
-
-
-def _record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
-    """Record the FQN of each compressed parameter on the simulator compressing it.
-
-    A simulator reading from anything other than a ``get_attr`` node keeps the
-    default name: an activation, or a weight arriving through a decompression op,
-    has no parameter to name.
-    """
-    simulators = dict(model.named_modules(remove_duplicate=False))
-    for node in model.graph.nodes:
-        if node.op != "call_module" or not node.args:
-            continue
-        simulator = simulators.get(str(node.target))
-        if isinstance(simulator, CompressionSimulatorBase) and node.args[0].op == "get_attr":
-            simulator.tensor_fqn = str(node.args[0].target)
diff --git a/src/coreai_opt/palettization/kmeans/palettizer.py b/src/coreai_opt/palettization/kmeans/palettizer.py
index 99215dbc..e0aa27d8 100644
--- a/src/coreai_opt/palettization/kmeans/palettizer.py
+++ b/src/coreai_opt/palettization/kmeans/palettizer.py
@@ -36,7 +36,6 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig, ModuleConfigDict
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_eager
 from coreai_opt.palettization.base_palettizer import _BasePalettizer
 from coreai_opt.palettization.config.palettization_config import (
     KMeansPalettizerConfig,
@@ -199,8 +198,6 @@ def prepare(
         logger.info("Preparing model for palettization")
         prepared_model = self._handler.prepare(self._model, example_inputs=example_inputs)
 
-        _record_tensor_fqns_eager(prepared_model)
-
         # Load precomputed sensitivities if provided
         if sensitivity_path is not None:
             logger.info(
diff --git a/src/coreai_opt/quantization/_eager/quantizer.py b/src/coreai_opt/quantization/_eager/quantizer.py
index 9e76dfbe..7f00a7f2 100644
--- a/src/coreai_opt/quantization/_eager/quantizer.py
+++ b/src/coreai_opt/quantization/_eager/quantizer.py
@@ -31,7 +31,6 @@
 from coreai_opt.config.compression_config import ModuleCompressionConfig
 from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.config.spec.base import CompressionSpec
-from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_eager
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -224,8 +223,6 @@ def _postprocess_prepared_model(model: nn.Module) -> None:
         Args:
             model (nn.Module): The model after eager prepare().
         """
-        _record_tensor_fqns_eager(model)
-
         _apply_weight_axis_defaults(model)
         validate_activation_axes(model)
 
diff --git a/src/coreai_opt/quantization/_graph/quantizer.py b/src/coreai_opt/quantization/_graph/quantizer.py
index 8974962a..44181ab5 100644
--- a/src/coreai_opt/quantization/_graph/quantizer.py
+++ b/src/coreai_opt/quantization/_graph/quantizer.py
@@ -46,7 +46,7 @@
 from coreai_opt._utils.version_utils import version_ge
 from coreai_opt.common import ExportBackend
 from coreai_opt.config.compression_config import ModuleConfigDict, _build_module_alias_map
-from coreai_opt.config.spec.compression_simulator import _record_tensor_fqns_graph
+from coreai_opt.config.spec.compression_simulator import CompressionSimulatorBase
 from coreai_opt.quantization._axis_defaults import (
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes,
@@ -100,6 +100,22 @@
 logger = logging.getLogger(__name__)
 
 
+def _record_tensor_fqns_graph(model: torch.fx.GraphModule) -> None:
+    """Record the FQN of each compressed parameter on the simulator compressing it.
+
+    A simulator reading from anything other than a ``get_attr`` node keeps the
+    default name: an activation, or a weight arriving through a decompression op,
+    has no parameter to name.
+    """
+    simulators = dict(model.named_modules(remove_duplicate=False))
+    for node in model.graph.nodes:
+        if node.op != "call_module" or not node.args:
+            continue
+        simulator = simulators.get(str(node.target))
+        if isinstance(simulator, CompressionSimulatorBase) and node.args[0].op == "get_attr":
+            simulator.tensor_fqn = str(node.args[0].target)
+
+
 class _OpConfigLevel(Enum):
     """
     Enum to specify the op-level config type within a module config.

From 105582196a00b99a02a6f793a28bb7a9f0a3c3c4 Mon Sep 17 00:00:00 2001
From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
Date: Wed, 2 Sep 2026 16:53:49 -0700
Subject: [PATCH 10/10] fix precommit

Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com>
---
 src/coreai_opt/config/spec/compression_simulator.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/coreai_opt/config/spec/compression_simulator.py b/src/coreai_opt/config/spec/compression_simulator.py
index 980f8b6c..d47465b7 100644
--- a/src/coreai_opt/config/spec/compression_simulator.py
+++ b/src/coreai_opt/config/spec/compression_simulator.py
@@ -46,4 +46,3 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
             flowing through
         """
         pass
-