From ddfefca596878b18804be740846317d13aa15e39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dr=20Chamyoung=20=E5=8C=BB=E8=80=85?= Date: Wed, 5 Aug 2026 03:19:38 +0530 Subject: [PATCH 1/3] fix(torch-utils): use .default overload for conv_transpose2d/3d compat with torch>=2.9 --- src/coreai_opt/_utils/torch_utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 475c180..58ea771 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/src/coreai_opt/_utils/torch_utils.py @@ -40,12 +40,19 @@ torch.ops.aten.conv2d.default: torch.nn.Conv2d, torch.ops.aten.conv3d.default: torch.nn.Conv3d, torch.ops.aten.conv_transpose1d.default: torch.nn.ConvTranspose1d, - torch.ops.aten.conv_transpose2d.input: torch.nn.ConvTranspose2d, - torch.ops.aten.conv_transpose3d.input: torch.nn.ConvTranspose3d, + torch.ops.aten.conv_transpose2d.default: torch.nn.ConvTranspose2d, + torch.ops.aten.conv_transpose3d.default: 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""" From bb9b5891d02dd90d12a19c9f67f6a0001a0246a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dr=20Chamyoung=20=E5=8C=BB=E8=80=85?= Date: Thu, 6 Aug 2026 07:48:34 +0530 Subject: [PATCH 2/3] test(torch-utils): add regression tests for conv_transpose2d/3d overload fix --- tests/test_utils/test_torch_utils.py | 62 ++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_utils/test_torch_utils.py b/tests/test_utils/test_torch_utils.py index 8e149ec..1ee3f3d 100644 --- a/tests/test_utils/test_torch_utils.py +++ b/tests/test_utils/test_torch_utils.py @@ -13,11 +13,13 @@ 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: @@ -163,3 +165,63 @@ 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 From a761ed028b6fd55e5e5535e516c40996fe7fcb3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dr=20Chamyoung=20=E5=8C=BB=E8=80=85?= Date: Fri, 7 Aug 2026 06:03:56 +0530 Subject: [PATCH 3/3] test(axis-defaults): add graph-mode regression tests for ConvTranspose2d/3d axis resolution --- src/coreai_opt/_utils/torch_utils.py | 11 +-- tests/quantization/test_axis_defaults.py | 89 ++++++++++++++++++++++++ tests/test_utils/test_torch_utils.py | 62 ----------------- 3 files changed, 91 insertions(+), 71 deletions(-) diff --git a/src/coreai_opt/_utils/torch_utils.py b/src/coreai_opt/_utils/torch_utils.py index 58ea771..475c180 100644 --- a/src/coreai_opt/_utils/torch_utils.py +++ b/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""" diff --git a/tests/quantization/test_axis_defaults.py b/tests/quantization/test_axis_defaults.py index e09fd76..17fed07 100644 --- a/tests/quantization/test_axis_defaults.py +++ b/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 + diff --git a/tests/test_utils/test_torch_utils.py b/tests/test_utils/test_torch_utils.py index 1ee3f3d..8e149ec 100644 --- a/tests/test_utils/test_torch_utils.py +++ b/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