Arm backend: Fix transposed conv lowering

backends/arm/_passes/size_adjust_input_pass.py

@@ -63,9 +63,13 @@ def _greater_than(input: SymIntLike, other: int) -> bool | torch.SymBool:
 
 
 def get_slices_convolution(conv_node: torch.fx.Node) -> Slices:
-    slices = []
+    slices: Slices = []
 
-    input_node, weight, _, stride_hw, pad_hw, dilation_hw, _, _, _ = conv_node.args
+    input_node, weight, _, stride_hw, pad_hw, dilation_hw, transposed, _, _ = (
+        conv_node.args
+    )
+    if transposed:
+        return slices
     weight_shape = cast(torch.fx.Node, weight).meta["val"].shape
     input_shape = cast(torch.fx.Node, input_node).meta["val"].shape
     spatial_rank = len(input_shape) - 2

backends/arm/test/passes/test_size_adjust_input_pass.py

@@ -42,6 +42,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.conv(x)
 
 
+class TransposeConvModule(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.conv = torch.nn.ConvTranspose2d(
+            in_channels=3,
+            out_channels=6,
+            kernel_size=3,
+            stride=2,
+            output_padding=1,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
 def _needs_truncation(input_length, kernel_size, stride, padding):
     return _greater_than((input_length + 2 * padding - kernel_size) % stride, padding)
 
@@ -115,6 +130,30 @@ def test_size_adjust_input_static_conv_no_adjustment_needed():
     ), "No slice nodes should be inserted when no adjustment is needed"
 
 
+def test_size_adjust_input_skips_transpose_conv2d() -> None:
+    model = TransposeConvModule()
+    example_inputs = (torch.randn(1, 3, 16, 16),)
+    edge_model = to_edge(export(model, example_inputs))
+    edge_model = edge_model.transform([SizeAdjustInputPass()])
+    gm = edge_model.exported_program().graph_module
+
+    conv_node = next(
+        n
+        for n in gm.graph.nodes
+        if n.op == "call_function"
+        and n.target == exir_ops.edge.aten.convolution.default
+    )
+    input_node = conv_node.args[0]
+    assert input_node.meta["val"].shape == example_inputs[0].shape
+
+    slice_nodes = [
+        n
+        for n in gm.graph.nodes
+        if n.op == "call_function" and n.target == exir_ops.edge.aten.slice_copy.Tensor
+    ]
+    assert len(slice_nodes) == 0
+
+
 def test_size_adjust_input_dynamic_conv2d():
     kernel_size, stride, padding = 3, 3, 1
     model = ConvModule(kernel_size=kernel_size, stride=stride, padding=padding)