Fix Conv3d w8a8 quantization rank mismatch

coremltools/converters/mil/frontend/torch/converter.py

@@ -947,12 +947,14 @@ def _construct_quantization_op(
         zero_point: Optional[np.ndarray] = None
         if compression_info.zero_point is not None:
             zero_point = compression_info.zero_point.detach().numpy()
-        # For conv/conv_transpose, the weight has rank=4, so we auto-expand scale and zero-point if
-        # it only has two elements.
-        if len(weight.shape) == 4 and len(scale.shape) == 2:
-            scale = np.expand_dims(np.expand_dims(scale, axis=-1), axis=-1)
-            if zero_point is not None:
-                zero_point = np.expand_dims(np.expand_dims(zero_point, axis=-1), axis=-1)
+        # For conv/conv_transpose, the weight has rank >= 4 (rank=4 for Conv2d, rank=5 for Conv3d),
+        # so we auto-expand scale and zero-point if it only has two elements.
+        if len(weight.shape) >= 4 and len(scale.shape) == 2:
+            n_expand = len(weight.shape) - len(scale.shape)
+            for _ in range(n_expand):
+                scale = np.expand_dims(scale, axis=-1)
+                if zero_point is not None:
+                    zero_point = np.expand_dims(zero_point, axis=-1)
 
         if compressed_var is not None and compressed_var.op.op_type == "constexpr_lut_to_dense":
             # The quantization on lut could lead to extra two dims at the end.

coremltools/test/optimize/torch/conversion/quantization/test_quantization_conversion.py

@@ -4,6 +4,8 @@
 #  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 
 import pytest
+import torch
+import torch.nn as nn
 
 ct = pytest.importorskip("coremltools")
 import coremltools.test.optimize.torch.conversion.conversion_utils as util
@@ -60,6 +62,42 @@ def test_linear_quantizer(config, model, mnist_example_input, request):
 # endregion
 
 
+# region LinearQuantizer Conv3d
+def test_linear_quantizer_conv3d_w8a8():
+    """Regression test for https://github.com/apple/coremltools/issues/2452.
+
+    w8a8 symmetric quantization of Conv3d failed with:
+        ValueError: the `weight` should have same rank as `scale`,
+        but got (1, 1, 3, 3, 3) vs (1, 1)
+    because the per-tensor scale was not reshaped to match the 5-D weight rank.
+    """
+
+    class Conv3dModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv1 = nn.Conv3d(1, 1, kernel_size=3, padding=1)
+
+        def forward(self, input_1):
+            return self.conv1(input_1)
+
+    example_input = torch.randn(1, 1, 8, 8, 8)
+    config = LinearQuantizerConfig.from_dict(
+        {"global_config": {"quantization_scheme": "symmetric"}}
+    )
+    quantizer = LinearQuantizer(Conv3dModel().eval(), config)
+    quantized_model = get_quantized_model(quantizer, example_input)
+
+    util.convert_and_verify(
+        quantized_model,
+        example_input,
+        minimum_deployment_target=ct.target.iOS17,
+        expected_ops=["constexpr_affine_dequantize"],
+    )
+
+
+# endregion
+
+
 # region GPTQ
 @pytest.mark.parametrize(
     "config",