fix(quantization): address review feedback on weight FQ calibration change

- Move enable_weight_fake_quant/disable_activation_fake_quant into a new
  _fake_quant_utils module so they can import FakeQuantizeImplBase and
  CompressionTargetTensor at module scope instead of via local imports,
  keeping _utils.py free of dependencies on other quantization modules.
- Use basic_config (weight + activation quantization) instead of
  input_activation_only_config in test_calibration_mode so the weight-FQ
  branch is actually exercised; only assert scale drift for activation
  modules, since weight ranges are fixed at prepare time.
- Drop the changelog.d entry.

Author: Utkarsh Simha

changelog.d/weight-fq-calibration-mode.changed

@@ -39,7 +39,7 @@
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes as _validate_activation_axes,
 )
-from coreai_opt.quantization._utils import (
+from coreai_opt.quantization._fake_quant_utils import (
     disable_activation_fake_quant as _disable_activation_fake_quant,
     enable_weight_fake_quant as _enable_weight_fake_quant,
 )

src/coreai_opt/quantization/_eager/quantizer.py

@@ -0,0 +1,49 @@
+# 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
+
+"""Helpers for toggling fake quantization by quantization target."""
+
+import torch
+
+from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+
+
+def disable_activation_fake_quant(module: torch.nn.Module) -> None:
+    """Disable fake quantization on activation FakeQuantize modules only.
+
+    Mirrors ``torchao.quantization.pt2e.fake_quantize.disable_fake_quant`` but
+    skips weight FQ modules. Used by ``calibration_mode`` so activation
+    observers see the effect of quantized weights when collecting statistics.
+
+    Args:
+        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
+            module that is not a ``FakeQuantizeImplBase`` whose
+            ``quantization_target`` is ``ACTIVATION``.
+    """
+    if (
+        isinstance(module, FakeQuantizeImplBase)
+        and module.quantization_target == CompressionTargetTensor.ACTIVATION
+    ):
+        module.disable_fake_quant()
+
+
+def enable_weight_fake_quant(module: torch.nn.Module) -> None:
+    """Enable fake quantization on weight FakeQuantize modules only.
+
+    Mirrors ``torchao.quantization.pt2e.fake_quantize.enable_fake_quant`` but
+    skips activation FQ modules. Companion to
+    :func:`disable_activation_fake_quant` used by ``calibration_mode``.
+
+    Args:
+        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
+            module that is not a ``FakeQuantizeImplBase`` whose
+            ``quantization_target`` is ``WEIGHT``.
+    """
+    if (
+        isinstance(module, FakeQuantizeImplBase)
+        and module.quantization_target == CompressionTargetTensor.WEIGHT
+    ):
+        module.enable_fake_quant()

src/coreai_opt/quantization/_fake_quant_utils.py

@@ -60,7 +60,7 @@
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes as _validate_activation_axes,
 )
-from coreai_opt.quantization._utils import (
+from coreai_opt.quantization._fake_quant_utils import (
     disable_activation_fake_quant as _disable_activation_fake_quant,
     enable_weight_fake_quant as _enable_weight_fake_quant,
 )

src/coreai_opt/quantization/_graph/quantizer.py

@@ -46,51 +46,3 @@ def get_quantization_shapes(
         reduced_shape[i] = 1
 
     return original_shape, blockwise_shape, reduced_shape
-
-
-def disable_activation_fake_quant(module: torch.nn.Module) -> None:
-    """Disable fake quantization on activation FakeQuantize modules only.
-
-    Mirrors ``torchao.quantization.pt2e.fake_quantize.disable_fake_quant`` but
-    skips weight FQ modules. Used by ``calibration_mode`` so activation
-    observers see the effect of quantized weights when collecting statistics.
-
-    Args:
-        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
-            module that is not a ``FakeQuantizeImplBase`` whose
-            ``quantization_target`` is ``ACTIVATION``.
-    """
-    # Lazy imports break a cycle: spec.fake_quantize imports get_quantization_shapes
-    # from this module.
-    from coreai_opt.config.spec import CompressionTargetTensor  # noqa: PLC0415
-    from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase  # noqa: PLC0415
-
-    if (
-        isinstance(module, FakeQuantizeImplBase)
-        and module.quantization_target == CompressionTargetTensor.ACTIVATION
-    ):
-        module.disable_fake_quant()
-
-
-def enable_weight_fake_quant(module: torch.nn.Module) -> None:
-    """Enable fake quantization on weight FakeQuantize modules only.
-
-    Mirrors ``torchao.quantization.pt2e.fake_quantize.enable_fake_quant`` but
-    skips activation FQ modules. Companion to
-    :func:`disable_activation_fake_quant` used by ``calibration_mode``.
-
-    Args:
-        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
-            module that is not a ``FakeQuantizeImplBase`` whose
-            ``quantization_target`` is ``WEIGHT``.
-    """
-    # Lazy imports break a cycle: spec.fake_quantize imports get_quantization_shapes
-    # from this module.
-    from coreai_opt.config.spec import CompressionTargetTensor  # noqa: PLC0415
-    from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase  # noqa: PLC0415
-
-    if (
-        isinstance(module, FakeQuantizeImplBase)
-        and module.quantization_target == CompressionTargetTensor.WEIGHT
-    ):
-        module.enable_fake_quant()

src/coreai_opt/quantization/_utils.py

@@ -530,23 +530,30 @@ def test_finalize_state_dict_safetensors_roundtrip(self, basic_config, tmp_path)
 
         assert torch.equal(out_before_roundtrip, out_after_roundtrip)
 
-    def test_calibration_mode(self, simple_model, input_activation_only_config, example_input):
+    def test_calibration_mode(self, simple_model, basic_config, example_input):
         """
         Test that calibration mode works as expected, and scales are getting updated
         """
-        quantizer = Quantizer(simple_model, input_activation_only_config)
+        quantizer = Quantizer(simple_model, basic_config)
         simple_model.eval()
         prepared_model = quantizer.prepare((example_input,))
 
         fake_quant_modules = [
             m for m in prepared_model.modules() if isinstance(m, FakeQuantizeImplBase)
         ]
+        activation_fake_quant_modules = [
+            m
+            for m in fake_quant_modules
+            if m.quantization_target == CompressionTargetTensor.ACTIVATION
+        ]
 
         for module in fake_quant_modules:
             assert module.observer_enabled.item() == 0
             assert module.fake_quant_enabled.item() == 1
 
-        pre_calibration_scales = [mod.calculate_qparams()[0].clone() for mod in fake_quant_modules]
+        pre_calibration_scales = [
+            mod.calculate_qparams()[0].clone() for mod in activation_fake_quant_modules
+        ]
 
         with quantizer.calibration_mode():
             simple_model.eval()
@@ -563,8 +570,12 @@ def test_calibration_mode(self, simple_model, input_activation_only_config, exam
                 )
                 assert module.fake_quant_enabled.item() == expected_fq
 
-        post_calibration_scales = [mod.calculate_qparams()[0].clone() for mod in fake_quant_modules]
+        post_calibration_scales = [
+            mod.calculate_qparams()[0].clone() for mod in activation_fake_quant_modules
+        ]
 
+        # Only activation scales are expected to move here: weight ranges are
+        # fixed at prepare time and don't depend on calibration data.
         for pre_scale, post_scale in zip(
             pre_calibration_scales, post_calibration_scales, strict=True
         ):