test(axis-defaults): add graph-mode regression tests for ConvTranspose2d/3d axis resolution

Author: InboraStudio

src/coreai_opt/_utils/torch_utils.py

@@ -40,19 +40,12 @@
     torch.ops.aten.conv2d.default: torch.nn.Conv2d,
     torch.ops.aten.conv3d.default: torch.nn.Conv3d,
     torch.ops.aten.conv_transpose1d.default: torch.nn.ConvTranspose1d,
-    torch.ops.aten.conv_transpose2d.default: torch.nn.ConvTranspose2d,
-    torch.ops.aten.conv_transpose3d.default: torch.nn.ConvTranspose3d,
+    torch.ops.aten.conv_transpose2d.input: torch.nn.ConvTranspose2d,
+    torch.ops.aten.conv_transpose3d.input: torch.nn.ConvTranspose3d,
     torch.ops.aten.linear.default: torch.nn.Linear,
     torch.ops.aten.embedding.default: torch.nn.Embedding,
 }
 
-# Backward-compat aliases for torch < 2.9
-try:
-    ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose2d.input] = torch.nn.ConvTranspose2d
-    ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose3d.input] = torch.nn.ConvTranspose3d
-except AttributeError:
-    pass
-
 
 class NamedModule(NamedTuple):
     """NamedTuple for holding name and module info"""

tests/quantization/test_axis_defaults.py

@@ -422,3 +422,92 @@ def test_different_fqs_independent(self):
         _apply_defaults(fq_map)
         assert fq_linear.granularity.axis == 0
         assert fq_conv_t.granularity.axis == 0
+
+
+class TestConvTransposeAxisDefaultsGraph:
+    """Graph-mode regression tests for ConvTranspose2d/3d axis default resolution.
+
+    ``ATEN_OP_TO_MODULE_TYPE`` maps the aten op emitted in the exported graph
+    (e.g. ``aten.conv_transpose2d.input``) to the corresponding ``nn.Module``
+    type so that ``_apply_defaults`` can look up the correct weight axis from
+    ``_WEIGHT_AXIS_SPECS``.  These tests verify the full coreai-opt path:
+    Quantizer.prepare → axis-defaults pass → correct axis on weight FQ.
+
+    ConvTranspose weight layout is ``[in_ch, out_ch, ...]``, so:
+      - per-channel axis (output channels) = 1
+      - per-block axis  (input channels)   = 0
+    """
+
+    @pytest.mark.parametrize(
+        ("make_model", "make_input"),
+        [
+            pytest.param(
+                lambda: nn.ConvTranspose2d(16, 8, 3, padding=1),
+                lambda: torch.randn(1, 16, 8, 8),
+                id="conv_transpose2d",
+            ),
+            pytest.param(
+                lambda: nn.ConvTranspose3d(16, 8, 3, padding=1),
+                lambda: torch.randn(1, 16, 4, 4, 4),
+                id="conv_transpose3d",
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        ("granularity", "expected_axis"),
+        [
+            pytest.param(PerChannelGranularity(axis=None), 1, id="per_channel_axis_1"),
+            pytest.param(
+                PerBlockGranularity(axis=None, block_size=_TEST_BLOCK_SIZE), 0, id="per_block_axis_0"
+            ),
+        ],
+    )
+    def test_axis_none_resolves_for_conv_transpose(
+        self,
+        make_model,
+        make_input,
+        granularity,
+        expected_axis,
+    ):
+        """ConvTranspose axis=None resolves to the correct default in graph mode.
+
+        Per-channel should resolve to axis 1 (output channels), per-block to
+        axis 0 (input channels), reflecting the [in_ch, out_ch, ...] weight layout.
+        """
+        config = _make_config(granularity, execution_mode="graph")
+        prepared = Quantizer(make_model(), config).prepare((make_input(),))
+
+        weight_fqs = _get_weight_fqs(prepared)
+        assert len(weight_fqs) == 1
+        assert weight_fqs[0].granularity.axis == expected_axis
+
+    @pytest.mark.parametrize(
+        ("make_model", "make_input"),
+        [
+            pytest.param(
+                lambda: nn.ConvTranspose2d(16, 8, 3, padding=1),
+                lambda: torch.randn(1, 16, 8, 8),
+                id="conv_transpose2d",
+            ),
+            pytest.param(
+                lambda: nn.ConvTranspose3d(16, 8, 3, padding=1),
+                lambda: torch.randn(1, 16, 4, 4, 4),
+                id="conv_transpose3d",
+            ),
+        ],
+    )
+    def test_prepare_calibrate_finalize_conv_transpose_graph(self, make_model, make_input):
+        """Full graph-mode workflow succeeds for ConvTranspose with axis=None."""
+        config = _make_config(PerChannelGranularity(axis=None), execution_mode="graph")
+        quantizer = Quantizer(make_model(), config)
+        example_input = make_input()
+
+        prepared = quantizer.prepare((example_input,))
+        with quantizer.calibration_mode():
+            prepared(example_input)
+        quantizer.finalize()
+
+        weight_fqs = _get_weight_fqs(prepared)
+        assert len(weight_fqs) == 1
+        assert weight_fqs[0].granularity.axis is not None
+

tests/test_utils/test_torch_utils.py

@@ -13,13 +13,11 @@
 
 from coreai_opt._utils.fx_utils import normalize_module_fqn
 from coreai_opt._utils.torch_utils import (
-    ATEN_OP_TO_MODULE_TYPE,
     mmap_module_state_dict,
     move_model_to_eval,
     move_model_to_train,
     normalize_axis,
 )
-from coreai_opt._utils.version_utils import version_ge
 
 
 class TestMoveModelContextManagers:
@@ -165,63 +163,3 @@ def test_raises_on_non_cpu_tensor(tmp_path):
 
         with pytest.raises(ValueError, match="requires CPU tensors"):
             mmap_module_state_dict(model, tmp_path / "model.safetensors")
-
-
-class TestAtenOpToModuleType:
-    """Tests for ATEN_OP_TO_MODULE_TYPE overload correctness across torch versions.
-
-    Regression test for: conv_transpose2d/3d were mapped to the `.input` overload
-    which was deprecated in PyTorch 2.9. The mapping must use `.default` as the
-    primary key while retaining the `.input` alias for backward compat on < 2.9.
-    """
-
-    @staticmethod
-    def test_conv_transpose2d_default_overload_present():
-        """conv_transpose2d.default must always be in the mapping (torch >= 2.9 primary path)."""
-        assert torch.ops.aten.conv_transpose2d.default in ATEN_OP_TO_MODULE_TYPE
-        assert ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose2d.default] is nn.ConvTranspose2d
-
-    @staticmethod
-    def test_conv_transpose3d_default_overload_present():
-        """conv_transpose3d.default must always be in the mapping (torch >= 2.9 primary path)."""
-        assert torch.ops.aten.conv_transpose3d.default in ATEN_OP_TO_MODULE_TYPE
-        assert ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose3d.default] is nn.ConvTranspose3d
-
-    @staticmethod
-    def test_conv_transpose2d_input_overload_compat():
-        """conv_transpose2d.input alias is present on torch < 2.9 for backward compat.
-
-        On torch >= 2.9 the `.input` overload may not exist at all; the test is
-        skipped in that case since the `.default` path fully covers those versions.
-        """
-        if not version_ge(torch, "2.9"):
-            # On < 2.9 the .input overload exists and must also resolve to ConvTranspose2d
-            assert torch.ops.aten.conv_transpose2d.input in ATEN_OP_TO_MODULE_TYPE
-            assert ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose2d.input] is nn.ConvTranspose2d
-        else:
-            # On >= 2.9 the .input overload may be absent; that's expected and fine
-            pytest.skip("conv_transpose2d.input overload not expected on torch >= 2.9")
-
-    @staticmethod
-    def test_conv_transpose3d_input_overload_compat():
-        """conv_transpose3d.input alias is present on torch < 2.9 for backward compat."""
-        if not version_ge(torch, "2.9"):
-            assert torch.ops.aten.conv_transpose3d.input in ATEN_OP_TO_MODULE_TYPE
-            assert ATEN_OP_TO_MODULE_TYPE[torch.ops.aten.conv_transpose3d.input] is nn.ConvTranspose3d
-        else:
-            pytest.skip("conv_transpose3d.input overload not expected on torch >= 2.9")
-
-    @staticmethod
-    def test_all_other_ops_unaffected():
-        """Sanity-check that unrelated ops in the mapping were not disturbed."""
-        expected = {
-            torch.ops.aten.conv1d.default: nn.Conv1d,
-            torch.ops.aten.conv2d.default: nn.Conv2d,
-            torch.ops.aten.conv3d.default: nn.Conv3d,
-            torch.ops.aten.conv_transpose1d.default: nn.ConvTranspose1d,
-            torch.ops.aten.linear.default: nn.Linear,
-            torch.ops.aten.embedding.default: nn.Embedding,
-        }
-        for op, module_type in expected.items():
-            assert op in ATEN_OP_TO_MODULE_TYPE
-            assert ATEN_OP_TO_MODULE_TYPE[op] is module_type