diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 3a38760..cbad1d2 100644 --- a/coreai_torch/_aten_to_core.py +++ b/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 = ( diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index 7b60157..4c24e96 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -1320,207 +1320,347 @@ def forward(self, x: Tensor) -> Tensor: ) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize( - "in_ch,out_ch,kernel,stride,padding,dilation,groups,bias,input_len", - [ - # Basic Conv1d, no padding - (3, 16, 3, 1, 0, 1, 1, True, 8), - # Stride=5, no padding, no bias - (3, 16, 3, 5, 0, 1, 1, False, 25), - # Non-zero padding (exercises the pad branch in replace_conv for rank==3) - (3, 16, 3, 1, 1, 1, 1, True, 8), - (4, 8, 5, 1, 2, 1, 1, False, 16), - # Padding + dilation - (3, 8, 3, 1, 2, 2, 1, True, 16), - ], -) -async def test_conv1d( - in_ch: int, - out_ch: int, - kernel: int, - stride: int, - padding: int, - dilation: int, - groups: int, - bias: bool, - input_len: int, - dynamic: bool, -) -> None: - x = torch.rand(2, in_ch, input_len) +class TestConvolution: + """Regular and transposed convolution conversions.""" - class Conv1dModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = nn.Conv1d( - in_ch, - out_ch, - kernel_size=kernel, - stride=stride, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - ) + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize( + "in_ch,out_ch,kernel,stride,padding,dilation,groups,bias,input_len", + [ + # Basic Conv1d, no padding + (3, 16, 3, 1, 0, 1, 1, True, 8), + # Stride=5, no padding, no bias + (3, 16, 3, 5, 0, 1, 1, False, 25), + # Non-zero padding (exercises the pad branch in replace_conv for rank==3) + (3, 16, 3, 1, 1, 1, 1, True, 8), + (4, 8, 5, 1, 2, 1, 1, False, 16), + # Padding + dilation + (3, 8, 3, 1, 2, 2, 1, True, 16), + ], + ) + async def test_conv1d( + self, + in_ch: int, + out_ch: int, + kernel: int, + stride: int, + padding: int, + dilation: int, + groups: int, + bias: bool, + input_len: int, + dynamic: bool, + ) -> None: + x = torch.rand(2, in_ch, input_len) - def forward(self, x: Tensor) -> Tensor: - return self.conv(x) + class Conv1dModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d( + in_ch, + out_ch, + kernel_size=kernel, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ) - model = Conv1dModel().eval() - # Batch (dim 0) and spatial length (dim 2) can be dynamic; channel (dim 1) is fixed by the layer. - # min_len ensures output_len >= 2 (avoids specialization when output == 1) and length >= kernel. - # max=2**20 keeps us within the upper bound torch.export derives from conv arithmetic. - if dynamic: - min_len = max(kernel, stride + dilation * (kernel - 1) + 1 - 2 * padding) - dynamic_shapes = { - "x": { - 0: torch.export.Dim("batch", min=1), - 2: torch.export.Dim("length", min=min_len, max=2**20), - } - } - else: - dynamic_shapes = None - await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + def forward(self, x: Tensor) -> Tensor: + return self.conv(x) + model = Conv1dModel().eval() + # Batch (dim 0) and spatial length (dim 2) can be dynamic; channel (dim 1) is fixed by the layer. + # min_len ensures output_len >= 2 (avoids specialization when output == 1) and length >= kernel. + # max=2**20 keeps us within the upper bound torch.export derives from conv arithmetic. + if dynamic: + min_len = max(kernel, stride + dilation * (kernel - 1) + 1 - 2 * padding) + dynamic_shapes = { + "x": { + 0: torch.export.Dim("batch", min=1), + 2: torch.export.Dim("length", min=min_len, max=2**20), + } + } + else: + dynamic_shapes = None + await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize( - "in_ch,out_ch,kernel,padding,dilation,groups,bias,input_shape", - [ - # Basic 2D conv with bias - (3, 16, 3, 1, 1, 1, True, (1, 3, 8, 8)), - # Without bias - (3, 16, 3, 1, 1, 1, False, (1, 3, 8, 8)), - # Depthwise: groups == in_channels - (4, 4, 3, 1, 1, 4, True, (1, 4, 8, 8)), - (8, 8, 3, 1, 1, 8, True, (2, 8, 6, 6)), - # Dilation > 1 (atrous convolution) - (3, 8, 3, 2, 2, 1, True, (1, 3, 16, 16)), - (3, 8, 3, (2, 3), (2, 3), 1, True, (1, 3, 16, 16)), - ], -) -async def test_conv2d( - in_ch: int, - out_ch: int, - kernel: int, - padding: int | tuple[int, int], - dilation: int | tuple[int, int], - groups: int, - bias: bool, - input_shape: tuple[int, int, int, int], - dynamic: bool, -) -> None: - x = torch.rand(2, *input_shape[1:]) + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize( + "in_ch,out_ch,kernel,padding,dilation,groups,bias,input_shape", + [ + # Basic 2D conv with bias + (3, 16, 3, 1, 1, 1, True, (1, 3, 8, 8)), + # Without bias + (3, 16, 3, 1, 1, 1, False, (1, 3, 8, 8)), + # Depthwise: groups == in_channels + (4, 4, 3, 1, 1, 4, True, (1, 4, 8, 8)), + (8, 8, 3, 1, 1, 8, True, (2, 8, 6, 6)), + # Dilation > 1 (atrous convolution) + (3, 8, 3, 2, 2, 1, True, (1, 3, 16, 16)), + (3, 8, 3, (2, 3), (2, 3), 1, True, (1, 3, 16, 16)), + ], + ) + async def test_conv2d( + self, + in_ch: int, + out_ch: int, + kernel: int, + padding: int | tuple[int, int], + dilation: int | tuple[int, int], + groups: int, + bias: bool, + input_shape: tuple[int, int, int, int], + dynamic: bool, + ) -> None: + x = torch.rand(2, *input_shape[1:]) - class Conv2dModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = nn.Conv2d( - in_ch, - out_ch, - kernel_size=kernel, - padding=padding, - dilation=dilation, - groups=groups, - bias=bias, - ) + class Conv2dModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_ch, + out_ch, + kernel_size=kernel, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ) - def forward(self, x: Tensor) -> Tensor: - return self.conv(x) + def forward(self, x: Tensor) -> Tensor: + return self.conv(x) - model = Conv2dModel().eval() - # Batch (dim 0) and spatial H/W (dims 2, 3) can be dynamic; channel (dim 1) is fixed by the layer. - dynamic_shapes = ( - { - "x": { - 0: torch.export.Dim("batch", min=1), - 2: torch.export.Dim("H", min=kernel), - 3: torch.export.Dim("W", min=kernel), + model = Conv2dModel().eval() + # Batch (dim 0) and spatial H/W (dims 2, 3) can be dynamic; channel (dim 1) is fixed by the layer. + dynamic_shapes = ( + { + "x": { + 0: torch.export.Dim("batch", min=1), + 2: torch.export.Dim("H", min=kernel), + 3: torch.export.Dim("W", min=kernel), + } } - } - if dynamic - else None + if dynamic + else None + ) + await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize( + "in_ch,out_ch,kernel,stride,padding,dilation,groups,bias,input_shape", + [ + # Basic 3D conv with bias, no padding + (3, 8, 3, 1, 0, 1, 1, True, (2, 3, 6, 6, 6)), + # Without bias + (3, 8, 3, 1, 0, 1, 1, False, (2, 3, 6, 6, 6)), + # Non-zero padding (exercises the pad branch in replace_conv for rank==5) + (3, 8, 3, 1, 1, 1, 1, True, (2, 3, 6, 6, 6)), + # Anisotropic kernel / padding / dilation across D, H, W + (3, 8, (3, 3, 3), 1, (1, 2, 0), (1, 2, 1), 1, True, (2, 3, 6, 8, 6)), + # Stride > 1 + (3, 8, 3, 2, 1, 1, 1, True, (2, 3, 8, 8, 8)), + # Depthwise: groups == in_channels + (4, 4, 3, 1, 1, 1, 4, True, (2, 4, 6, 6, 6)), + ], ) - await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + async def test_conv3d( + self, + in_ch: int, + out_ch: int, + kernel: int | tuple[int, int, int], + stride: int, + padding: int | tuple[int, int, int], + dilation: int | tuple[int, int, int], + groups: int, + bias: bool, + input_shape: tuple[int, int, int, int, int], + dynamic: bool, + ) -> None: + x = torch.rand(*input_shape) + class Conv3dModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv3d( + in_ch, + out_ch, + kernel_size=kernel, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize( - "in_ch,out_ch,kernel,stride,padding,dilation,groups,bias,input_shape", - [ - # Basic 3D conv with bias, no padding - (3, 8, 3, 1, 0, 1, 1, True, (2, 3, 6, 6, 6)), - # Without bias - (3, 8, 3, 1, 0, 1, 1, False, (2, 3, 6, 6, 6)), - # Non-zero padding (exercises the pad branch in replace_conv for rank==5) - (3, 8, 3, 1, 1, 1, 1, True, (2, 3, 6, 6, 6)), - # Anisotropic kernel / padding / dilation across D, H, W - (3, 8, (3, 3, 3), 1, (1, 2, 0), (1, 2, 1), 1, True, (2, 3, 6, 8, 6)), - # Stride > 1 - (3, 8, 3, 2, 1, 1, 1, True, (2, 3, 8, 8, 8)), - # Depthwise: groups == in_channels - (4, 4, 3, 1, 1, 1, 4, True, (2, 4, 6, 6, 6)), - ], -) -async def test_conv3d( - in_ch: int, - out_ch: int, - kernel: int | tuple[int, int, int], - stride: int, - padding: int | tuple[int, int, int], - dilation: int | tuple[int, int, int], - groups: int, - bias: bool, - input_shape: tuple[int, int, int, int, int], - dynamic: bool, -) -> None: - x = torch.rand(*input_shape) + def forward(self, x: Tensor) -> Tensor: + return self.conv(x) + + model = Conv3dModel().eval() + + # Batch (dim 0) and spatial D/H/W (dims 2, 3, 4) can be dynamic; channel (dim 1) is fixed. + # Per-axis min matches the conv1d logic: large enough that kernel fits AND output >= 2 + # (an output of 1 gets specialized, breaking the dynamic guard). + def _triple(v: int | tuple[int, int, int]) -> tuple[int, int, int]: + return v if isinstance(v, tuple) else (v, v, v) + + k = _triple(kernel) + p = _triple(padding) + d = _triple(dilation) + s = _triple(stride) + min_d = max(k[0], s[0] + d[0] * (k[0] - 1) + 1 - 2 * p[0]) + min_h = max(k[1], s[1] + d[1] * (k[1] - 1) + 1 - 2 * p[1]) + min_w = max(k[2], s[2] + d[2] * (k[2] - 1) + 1 - 2 * p[2]) + dynamic_shapes = ( + { + "x": { + 0: torch.export.Dim("batch", min=1), + 2: torch.export.Dim("D", min=min_d, max=2**20), + 3: torch.export.Dim("H", min=min_h, max=2**20), + 4: torch.export.Dim("W", min=min_w, max=2**20), + } + } + if dynamic + else None + ) + await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) - class Conv3dModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv = nn.Conv3d( - in_ch, - out_ch, - kernel_size=kernel, + @pytest.mark.parametrize("dynamic", [False, True]) + @pytest.mark.parametrize( + "in_channels,out_channels,kernel_size,stride,padding,dilation,output_padding,groups,is_1d", + [ + # ConvTranspose2d: Basic cases with stride=2 (upsampling) + (1, 1, 3, (2, 2), (1, 1), (1, 1), (0, 0), 1, False), + # ConvTranspose2d: Different kernel/stride combinations + (3, 6, 3, (1, 1), (0, 0), (1, 1), (0, 0), 1, False), + # ConvTranspose2d: With padding and stride + (2, 4, 3, (2, 2), (1, 1), (1, 1), (0, 0), 1, False), + # ConvTranspose2d: With dilation + (2, 4, 3, (1, 1), (2, 2), (2, 2), (0, 0), 1, False), + # ConvTranspose2d: With output_padding + (1, 1, 3, (2, 2), (1, 1), (1, 1), (1, 1), 1, False), + # ConvTranspose2d: Grouped convolution + (4, 4, 3, (2, 2), (1, 1), (1, 1), (0, 0), 2, False), + # ConvTranspose1d: Basic case + (1, 1, 3, (2,), (1,), (1,), (0,), 1, True), + # ConvTranspose1d: With stride and padding + (2, 4, 3, (2,), (1,), (1,), (0,), 1, True), + # ConvTranspose1d: With output_padding + (2, 4, 3, (2,), (1,), (1,), (1,), 1, True), + ], + ) + async def test_conv_transpose( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: Any, + padding: Any, + dilation: Any, + output_padding: Any, + groups: int, + is_1d: bool, + dynamic: bool, + ) -> None: + """Test conv_transpose1d and conv_transpose2d operations.""" + conv_layer: Any = None + if is_1d: + x = torch.rand(2, in_channels, 8) + conv_layer = nn.ConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride[0], + padding=padding[0], + dilation=dilation[0], + output_padding=output_padding[0], + groups=groups, + bias=True, + ) + else: + x = torch.rand(2, in_channels, 8, 8) + conv_layer = nn.ConvTranspose2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, + output_padding=output_padding, groups=groups, - bias=bias, + bias=True, ) - def forward(self, x: Tensor) -> Tensor: - return self.conv(x) - - model = Conv3dModel().eval() - - # Batch (dim 0) and spatial D/H/W (dims 2, 3, 4) can be dynamic; channel (dim 1) is fixed. - # Per-axis min matches the conv1d logic: large enough that kernel fits AND output >= 2 - # (an output of 1 gets specialized, breaking the dynamic guard). - def _triple(v: int | tuple[int, int, int]) -> tuple[int, int, int]: - return v if isinstance(v, tuple) else (v, v, v) - - k = _triple(kernel) - p = _triple(padding) - d = _triple(dilation) - s = _triple(stride) - min_d = max(k[0], s[0] + d[0] * (k[0] - 1) + 1 - 2 * p[0]) - min_h = max(k[1], s[1] + d[1] * (k[1] - 1) + 1 - 2 * p[1]) - min_w = max(k[2], s[2] + d[2] * (k[2] - 1) + 1 - 2 * p[2]) - dynamic_shapes = ( - { - "x": { - 0: torch.export.Dim("batch", min=1), - 2: torch.export.Dim("D", min=min_d, max=2**20), - 3: torch.export.Dim("H", min=min_h, max=2**20), - 4: torch.export.Dim("W", min=min_w, max=2**20), - } - } - if dynamic - else None + class ConvTransposeModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv_transpose = conv_layer + + def forward(self, x: Tensor) -> Tensor: + return self.conv_transpose(x) + + model = ConvTransposeModel().eval() + # Batch (dim 0) and spatial dimensions can be dynamic; channel (dim 1) is fixed by the layer. + if dynamic: + dynamic_shapes = ( + { + "x": { + 0: torch.export.Dim("batch", min=1), + 2: torch.export.Dim("L", min=kernel_size, max=2**20), + } + } + if is_1d + else { + "x": { + 0: torch.export.Dim("batch", min=1), + 2: torch.export.Dim("H", min=kernel_size), + 3: torch.export.Dim("W", min=kernel_size), + } + } + ) + else: + dynamic_shapes = None + 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), + ], ) - await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) + async def test_conv_transpose_static_shape_squeeze( + self, 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) class TestCopy: @@ -5565,103 +5705,6 @@ def forward(self, x: Tensor, row_idx: Tensor, col_idx: Tensor) -> Tensor: ) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize( - "in_channels,out_channels,kernel_size,stride,padding,dilation,output_padding,groups,is_1d", - [ - # ConvTranspose2d: Basic cases with stride=2 (upsampling) - (1, 1, 3, (2, 2), (1, 1), (1, 1), (0, 0), 1, False), - # ConvTranspose2d: Different kernel/stride combinations - (3, 6, 3, (1, 1), (0, 0), (1, 1), (0, 0), 1, False), - # ConvTranspose2d: With padding and stride - (2, 4, 3, (2, 2), (1, 1), (1, 1), (0, 0), 1, False), - # ConvTranspose2d: With dilation - (2, 4, 3, (1, 1), (2, 2), (2, 2), (0, 0), 1, False), - # ConvTranspose2d: With output_padding - (1, 1, 3, (2, 2), (1, 1), (1, 1), (1, 1), 1, False), - # ConvTranspose2d: Grouped convolution - (4, 4, 3, (2, 2), (1, 1), (1, 1), (0, 0), 2, False), - # ConvTranspose1d: Basic case - (1, 1, 3, (2,), (1,), (1,), (0,), 1, True), - # ConvTranspose1d: With stride and padding - (2, 4, 3, (2,), (1,), (1,), (0,), 1, True), - # ConvTranspose1d: With output_padding - (2, 4, 3, (2,), (1,), (1,), (1,), 1, True), - ], -) -async def test_conv_transpose( - in_channels: int, - out_channels: int, - kernel_size: int, - stride: Any, - padding: Any, - dilation: Any, - output_padding: Any, - groups: int, - is_1d: bool, - dynamic: bool, -) -> None: - """Test conv_transpose1d and conv_transpose2d operations.""" - conv_layer: Any = None - if is_1d: - x = torch.rand(2, in_channels, 8) - conv_layer = nn.ConvTranspose1d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride[0], - padding=padding[0], - dilation=dilation[0], - output_padding=output_padding[0], - groups=groups, - bias=True, - ) - else: - x = torch.rand(2, in_channels, 8, 8) - conv_layer = nn.ConvTranspose2d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, - output_padding=output_padding, - groups=groups, - bias=True, - ) - - class ConvTransposeModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.conv_transpose = conv_layer - - def forward(self, x: Tensor) -> Tensor: - return self.conv_transpose(x) - - model = ConvTransposeModel().eval() - # Batch (dim 0) and spatial dimensions can be dynamic; channel (dim 1) is fixed by the layer. - if dynamic: - dynamic_shapes = ( - { - "x": { - 0: torch.export.Dim("batch", min=1), - 2: torch.export.Dim("L", min=kernel_size, max=2**20), - } - } - if is_1d - else { - "x": { - 0: torch.export.Dim("batch", min=1), - 2: torch.export.Dim("H", min=kernel_size), - 3: torch.export.Dim("W", min=kernel_size), - } - } - ) - else: - dynamic_shapes = None - await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes) - - @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "input_shape,split_sizes,dim,dtype", diff --git a/tests/ops/test_ops_ir.py b/tests/ops/test_ops_ir.py index fc997d8..1e242cb 100644 --- a/tests/ops/test_ops_ir.py +++ b/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 - // 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) -> 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 + // 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) -> tensor<1x8x8x8xf32> + // CHECK-NEXT: coreai.output %[[R]] : tensor<1x8x8x8xf32> // CHECK-NEXT: } // CHECK-NEXT: } """,