Improve Functional convolution tests

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -1694,21 +1694,25 @@ def test_convolution1d(
         self,
         compute_unit,
         backend,
+        padding,
         width,
         in_channels,
         out_channels,
         kernel_size,
         stride,
-        padding,
         groups=1,
     ):
-        class DynamicConv(nn.Module):
+        if padding == "same" and stride != 1:
+            # configuration not supported
+            return
+
+        class FunctionalConv1D(nn.Module):
             def forward(self, input_data, weights):
                 return nn.functional.conv1d(
-                    input_data, weights, stride=stride, padding=padding
+                    input_data, weights, stride=stride, padding=padding, groups=groups
                 )
 
-        model = DynamicConv()
+        model = FunctionalConv1D()
         input_shape = [
             (1, in_channels, width),
             (out_channels, int(in_channels / groups), kernel_size),
@@ -1725,29 +1729,30 @@ def forward(self, input_data, weights):
             [
                 "compute_unit",
                 "backend",
+                "padding",
                 "height",
                 "width",
                 "in_channels",
                 "out_channels",
                 "kernel_size",
                 "stride",
-                "padding",
             ]
         ),
         [
-            (compute_unit, backend, *param)
-            for compute_unit, backend, param in itertools.product(
+            (compute_unit, backend, padding, *param)
+            for compute_unit, backend, padding, param in itertools.product(
                 compute_units,
                 backends,
+                ["same", "valid", 1, 0],
                 [
-                    (5, 3, 1, 1, 1, 2, 0),
-                    (3, 3, 1, 1, 1, 2, 1),
-                    (4, 3, 3, 3, 1, 2, 0),
-                    (7, 3, 3, 3, 1, 3, 0),
-                    (5, 5, 3, 3, 2, 1, 0),
-                    (3, 5, 3, 3, 1, 3, 0),
-                    (3, 5, 3, 3, 1, 3, 1),
-                    (7, 5, 3, 3, 2, 3, 1),
+                    (5, 3, 1, 1, 1, 2),
+                    (3, 3, 1, 1, 1, 2),
+                    (4, 3, 3, 3, 1, 2),
+                    (7, 3, 3, 3, 1, 3),
+                    (5, 5, 3, 3, 2, 1),
+                    (3, 5, 3, 3, 1, 3),
+                    (3, 5, 3, 3, 1, 3),
+                    (7, 5, 3, 3, 2, 3),
                 ],
             )
         ],
@@ -1756,31 +1761,118 @@ def test_convolution2d(
         self,
         compute_unit,
         backend,
+        padding,
         height,
         width,
         in_channels,
         out_channels,
         kernel_size,
         stride,
-        padding,
         groups=1,
     ):
-        class DynamicConv(nn.Module):
+        if padding == "same" and stride != 1:
+            # configuration not supported
+            return
+
+        class FunctionalConv2D(nn.Module):
             def forward(self, input_data, weights):
                 return nn.functional.conv2d(
-                    input_data, weights, stride=stride, padding=padding
+                    input_data, weights, stride=stride, padding=padding, groups=groups
                 )
 
-        model = DynamicConv()
+        model = FunctionalConv2D()
 
         input_shape = [
             (1, in_channels, height, width),
             (out_channels, int(in_channels / groups), kernel_size, kernel_size),
         ]
         self.run_compare_torch(
-            input_shape, model, backend=backend, compute_unit=compute_unit
+            input_shape,
+            model,
+            backend=backend,
+            compute_unit=compute_unit,
         )
 
+    @pytest.mark.parametrize(
+        ",".join(
+            [
+                "compute_unit",
+                "backend",
+                "padding",
+                "depth",
+                "height",
+                "width",
+                "in_channels",
+                "out_channels",
+                "kernel_size",
+                "stride",
+            ]
+        ),
+        [
+            (compute_unit, backend, padding, *param)
+            for compute_unit, backend, padding, param in itertools.product(
+                compute_units,
+                backends,
+                ["same", "valid", 1, 0],
+                [
+                    (5, 3, 2, 1, 1, 1, 2),
+                    (3, 3, 1, 1, 1, 1, 2),
+                    (4, 3, 3, 3, 3, 1, 2),
+                    (7, 3, 4, 3, 3, 1, 3),
+                    (5, 5, 3, 3, 3, 2, 1),
+                    (3, 5, 1, 3, 3, 1, 3),
+                    (3, 5, 4, 3, 3, 1, 3),
+                    (7, 5, 6, 3, 3, 2, 3),
+                ],
+            )
+        ],
+    )
+    def test_convolution3d(
+        self,
+        compute_unit,
+        backend,
+        padding,
+        depth,
+        height,
+        width,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        groups=1,
+    ):
+        if padding == "same" and stride != 1:
+            # configuration not supported
+            return
+
+        class FunctionalConv3D(nn.Module):
+            def forward(self, input_data, weights):
+                return nn.functional.conv3d(
+                    input_data, weights, stride=stride, padding=padding, groups=groups
+                )
+
+        model = FunctionalConv3D()
+        input_shape = [
+            (1, in_channels, depth, height, width),
+            (out_channels, int(in_channels / groups), kernel_size, kernel_size, kernel_size),
+        ]
+
+        if "neuralnetwork" in backend:
+            with pytest.raises(ValueError, match="3D Convolution doesn't support dynamic weights."):
+                self.run_compare_torch(
+                    input_shape,
+                    model,
+                    backend=backend,
+                    compute_unit=compute_unit,
+                )
+        else:
+            self.run_compare_torch(
+                input_shape,
+                model,
+                backend=backend,
+                compute_unit=compute_unit,
+            )
+
 
 class TestConvTranspose(TorchBaseTest):
     @pytest.mark.parametrize(