fix conv transpose lowering

Author: jakesabathia2

coreai_torch/_aten_to_core.py

@@ -1148,7 +1148,8 @@ def _conv_transpose(
     """Handles transposed convolution (conv_transpose1d and conv_transpose2d).
 
     For 1D, expands to 2D, performs conv_transpose2d, then shrinks back.
-    Handles output_padding via pre-padding input and post-cropping output.
+    ``padding`` and ``output_padding`` are handled natively by the Core AI
+    ``conv_transpose2d`` op (matching PyTorch semantics).
     """
     is_1d = x.type.rank == 3
     if is_1d:
@@ -1160,66 +1161,29 @@ def _conv_transpose(
         dilation = dilation + [1]
         output_padding = output_padding + [0]
 
-    x_rank = x.type.rank
-    effective_padding = padding
-    pre_pad_amt = [0] * (x_rank * 2)
-    post_crop_amt = [0] * (x_rank * 2)
-
-    if any(p > 0 for p in output_padding):
-        effective_padding = [0] * len(padding)
-        pre_pad_amt = [0] * (x_rank * 2)
-        post_crop_amt = [0] * (x_rank * 2)
-        # For each spatial dim: initialize symmetric crop from padding,
-        # then shift the output_padding amount from crop → pre-pad if needed
-        for i, (p, op) in enumerate(zip(padding, output_padding)):
-            before = 4 + 2 * i
-            after = 4 + 2 * i + 1
-            post_crop_amt[before] = p
-            post_crop_amt[after] = p
-            if post_crop_amt[after] >= op:
-                post_crop_amt[after] -= op
-            else:
-                pre_pad_amt[after] = op - post_crop_amt[after]
-                post_crop_amt[after] = 0
-
-    if any(p > 0 for p in pre_pad_amt):
-        x = coreai.pad(
-            x,
-            np.array(pre_pad_amt, dtype=np.uint32),
-            coreai.constant(0, dtype=x.type.element_type),
-        )
-    stride = coreai.constant(stride, np.uint32)
-    effective_padding = coreai.constant(effective_padding, np.uint32)
-    dilation = coreai.constant(dilation, np.uint32)
-    output_padding = coreai.constant([0, 0], dtype=np.uint32)
-    groups = coreai.constant(groups, np.uint32)
     result = coreai.conv_transpose2d(
         input=x,
         weight=weight,
-        stride=stride,
-        padding=effective_padding,
-        dilation=dilation,
-        output_pad=output_padding,
-        groups=groups,
+        stride=coreai.constant(stride, np.uint32),
+        padding=coreai.constant(padding, np.uint32),
+        dilation=coreai.constant(dilation, np.uint32),
+        output_pad=coreai.constant(output_padding, np.uint32),
+        groups=coreai.constant(groups, np.uint32),
     )
 
-    if any(p > 0 for p in post_crop_amt):
-        stop_val = coreai.sub(
-            coreai.cast(coreai.get_shape(result), dtype=np.int32),
-            [post_crop_amt[2 * d + 1] for d in range(x_rank)],
-        )
-        result = coreai.slice_(
-            result,
-            [post_crop_amt[2 * d] for d in range(x_rank)],
-            stop_val,
-            [1] * x_rank,
-        )
-
     if is_1d:
-        # Shrink back to 3D: [N,C,W,1] → [N,C,W]
-        result = coreai.reshape(
-            result, coreai.slice_(coreai.get_shape(result), [0], [3], [1])
-        )
+        # Shrink back to 3D: [N,C,W,1] → [N,C,W]. When the trailing (added) dim
+        # is statically 1, use shrink_dims — the inverse of the expand_dims
+        # above — which preserves the statically-known output shape so
+        # downstream ops (e.g. squeeze) stay static (rdar://181169322). Under a
+        # dynamic input the conv op reports every dim (incl. the added one) as
+        # dynamic, so fall back to a reshape driven by the runtime shape.
+        if result.type.shape[-1] == 1:
+            result = coreai.shrink_dims(result, [-1])
+        else:
+            result = coreai.reshape(
+                result, coreai.slice_(coreai.get_shape(result), [0], [3], [1])
+            )
 
     if bias is not None:
         bias_shape = (

tests/ops/test_ops.py

@@ -5662,6 +5662,46 @@ def forward(self, x: Tensor) -> Tensor:
     await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
 
 
+@pytest.mark.parametrize(
+    "is_1d, output_padding",
+    [
+        (True, 0),  # radar repro: conv_transpose1d + squeeze, static shapes
+        (True, 1),  # output_padding > padding (previously wrong output size)
+        (False, 0),
+        (False, 1),
+    ],
+)
+async def test_conv_transpose_static_shape_squeeze(
+    is_1d: bool, output_padding: int
+) -> None:
+    """A squeeze/reshape after a transposed conv with fully static input must
+    keep static shapes end-to-end so ``save_asset`` passes MLIR verification
+    (rdar://181169322: the 1D reshape-back used to erase the static shape,
+    tripping ``coreai.shrink_dims`` on dynamic dims)."""
+
+    class Model(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.ct = (
+                nn.ConvTranspose1d(4, 1, 8, stride=2, output_padding=output_padding)
+                if is_1d
+                else nn.ConvTranspose2d(
+                    4, 1, 3, stride=2, output_padding=output_padding
+                )
+            )
+
+        def forward(self, x: Tensor) -> Tensor:
+            # squeeze the singleton out-channel dim, then reshape — the pattern
+            # that previously produced `shrink_dims` on dynamic dims.
+            y = self.ct(x).squeeze(1)
+            return y.reshape(y.shape[0], -1)
+
+    x = torch.randn(1, 4, 32) if is_1d else torch.randn(1, 4, 8, 8)
+    # validate_numerical_output saves the asset (exercising MLIR verification)
+    # and checks numerics against torch eager.
+    await validate_numerical_output(model=Model().eval(), x=x)
+
+
 @pytest.mark.parametrize("dynamic", [False, True])
 @pytest.mark.parametrize(
     "input_shape,split_sizes,dim,dtype",

tests/ops/test_ops_ir.py

@@ -2614,17 +2614,12 @@ def forward(self, x: Tensor) -> Tensor:
             check_file="""
                 // CHECK-LABEL: module {
                 // CHECK-NEXT:   coreai.graph @main(%[[ARG0:.*]]: tensor<1x3x4x4xf32> {coreai.name = "x"}) -> (tensor<1x8x8x8xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
-                // CHECK-NEXT:     %[[V0:.*]] = coreai.constant dense<[1, 8, 9, 9]> : tensor<4xsi32>
-                // CHECK-NEXT:     %[[V1:.*]] = coreai.constant dense<1> : tensor<4xsi32>
-                // CHECK-NEXT:     %[[V2:.*]] = coreai.constant dense<[0, 0, 1, 1]> : tensor<4xsi32>
-                // CHECK-NEXT:     %[[V3:.*]] = coreai.constant dense<{{.*}}> : tensor<3x8x3x3xf32>
-                // CHECK-NEXT:     %[[V4:.*]] = coreai.constant dense<2> : tensor<2xui32>
-                // CHECK-NEXT:     %[[V5:.*]] = coreai.constant dense<0> : tensor<2xui32>
-                // CHECK-NEXT:     %[[V6:.*]] = coreai.constant dense<1> : tensor<2xui32>
-                // CHECK-NEXT:     %[[V7:.*]] = coreai.constant dense<1> : tensor<ui32>
-                // CHECK-NEXT:     %[[V8:.*]] = coreai.conv_transpose2d %[[ARG0]], %[[V3]], %[[V4]], %[[V5]], %[[V6]], %[[V5]], %[[V7]] : (tensor<1x3x4x4xf32>, tensor<3x8x3x3xf32>, tensor<2xui32>, tensor<2xui32>, tensor<2xui32>, tensor<2xui32>, tensor<ui32>) -> tensor<1x8x9x9xf32>
-                // CHECK-NEXT:     %[[V9:.*]] = coreai.slice %[[V8]], %[[V2]], %[[V0]], %[[V1]] : (tensor<1x8x9x9xf32>, tensor<4xsi32>, tensor<4xsi32>, tensor<4xsi32>) -> tensor<1x8x8x8xf32>
-                // CHECK-NEXT:     coreai.output %[[V9]] : tensor<1x8x8x8xf32>
+                // CHECK-NEXT:     %[[W:.*]] = coreai.constant dense<{{.*}}> : tensor<3x8x3x3xf32>
+                // CHECK-NEXT:     %[[STRIDE:.*]] = coreai.constant dense<2> : tensor<2xui32>
+                // CHECK-NEXT:     %[[ONE:.*]] = coreai.constant dense<1> : tensor<2xui32>
+                // CHECK-NEXT:     %[[GROUPS:.*]] = coreai.constant dense<1> : tensor<ui32>
+                // CHECK-NEXT:     %[[R:.*]] = coreai.conv_transpose2d %[[ARG0]], %[[W]], %[[STRIDE]], %[[ONE]], %[[ONE]], %[[ONE]], %[[GROUPS]] : (tensor<1x3x4x4xf32>, tensor<3x8x3x3xf32>, tensor<2xui32>, tensor<2xui32>, tensor<2xui32>, tensor<2xui32>, tensor<ui32>) -> tensor<1x8x8x8xf32>
+                // CHECK-NEXT:     coreai.output %[[R]] : tensor<1x8x8x8xf32>
                 // CHECK-NEXT:   }
                 // CHECK-NEXT: }
             """,