cleanup conv1d

Author: mshr-h

python/tvm/relax/frontend/torch/fx_translator.py

@@ -788,33 +788,52 @@ def _conv1d(self, node: fx.node.Node) -> relax.Var:
             groups=module.groups,
         )
 
-    def _conv3d(self, node: fx.node.Node) -> relax.Var:
+    def _conv1d_functional(self, node: fx.node.Node) -> relax.Var:
+        args = self.retrieve_args(node)
+        x = args[0]
+        weight = args[1]
+        bias = args[2] if len(args) > 2 else None
+        stride = args[3] if len(args) > 3 else 1
+        padding = args[4] if len(args) > 4 else 0
+        dilation = args[5] if len(args) > 5 else 1
+        groups = args[6] if len(args) > 6 else 1
+        return self._conv1d_impl(
+            x,
+            weight,
+            bias=bias,
+            strides=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+
+    def _conv1d_transpose(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
         weight = self.params[module.weight]
 
-        conv3d = self.block_builder.emit(
-            relax.op.nn.conv3d(
+        conv1d_transpose = self.block_builder.emit(
+            relax.op.nn.conv1d_transpose(
                 x,
                 weight,
                 strides=module.stride,
                 padding=module.padding,
                 dilation=module.dilation,
                 groups=module.groups,
-                data_layout="NCDHW",
-                kernel_layout="OIDHW",
+                data_layout="NCW",
+                kernel_layout="OIW",
                 out_dtype="float32",
             )
         )
 
         if module.bias is None:
-            return conv3d
+            return conv1d_transpose
 
         bias = self.params[module.bias]
         assert len(self.shape_of(bias)) == 1
-        bias = relax.op.reshape(bias, (1, -1, 1, 1, 1))
+        bias = relax.op.reshape(bias, (1, -1, 1))
 
-        return self.block_builder.emit(relax.op.add(conv3d, bias))
+        return self.block_builder.emit(relax.op.add(conv1d_transpose, bias))
 
     def _conv2d_impl(
         self,
@@ -846,7 +865,25 @@ def _conv2d_impl(
         bias = relax.op.reshape(bias, (1, -1, 1, 1))
         return self.block_builder.emit(relax.op.add(conv2d, bias))
 
-    def _conv1d_functional(self, node: fx.node.Node) -> relax.Var:
+    def _conv2d(self, node: fx.node.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        module = self.named_modules[node.target]
+        weight = self.params[module.weight]
+        bias = None
+        if module.bias is not None:
+            bias = self.params[module.bias]
+
+        return self._conv2d_impl(
+            x,
+            weight,
+            bias=bias,
+            strides=module.stride,
+            padding=module.padding,
+            dilation=module.dilation,
+            groups=module.groups,
+        )
+
+    def _conv2d_functional(self, node: fx.node.Node) -> relax.Var:
         args = self.retrieve_args(node)
         x = args[0]
         weight = args[1]
@@ -855,7 +892,7 @@ def _conv1d_functional(self, node: fx.node.Node) -> relax.Var:
         padding = args[4] if len(args) > 4 else 0
         dilation = args[5] if len(args) > 5 else 1
         groups = args[6] if len(args) > 6 else 1
-        return self._conv1d_impl(
+        return self._conv2d_impl(
             x,
             weight,
             bias=bias,
@@ -865,98 +902,61 @@ def _conv1d_functional(self, node: fx.node.Node) -> relax.Var:
             groups=groups,
         )
 
-    def _conv1d_transpose(self, node: fx.node.Node) -> relax.Var:
+    def _conv2d_transpose(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
         weight = self.params[module.weight]
 
-        conv1d_transpose = self.block_builder.emit(
-            relax.op.nn.conv1d_transpose(
+        conv2d_transpose = self.block_builder.emit(
+            relax.op.nn.conv2d_transpose(
                 x,
                 weight,
                 strides=module.stride,
                 padding=module.padding,
                 dilation=module.dilation,
                 groups=module.groups,
-                data_layout="NCW",
-                kernel_layout="OIW",
+                data_layout="NCHW",
+                kernel_layout="OIHW",
                 out_dtype="float32",
             )
         )
 
         if module.bias is None:
-            return conv1d_transpose
+            return conv2d_transpose
 
         bias = self.params[module.bias]
         assert len(self.shape_of(bias)) == 1
-        bias = relax.op.reshape(bias, (1, -1, 1))
+        bias = relax.op.reshape(bias, (1, -1, 1, 1))
 
-        return self.block_builder.emit(relax.op.add(conv1d_transpose, bias))
+        return self.block_builder.emit(relax.op.add(conv2d_transpose, bias))
 
-    def _conv2d_transpose(self, node: fx.node.Node) -> relax.Var:
+    def _conv3d(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
         weight = self.params[module.weight]
 
-        conv2d_transpose = self.block_builder.emit(
-            relax.op.nn.conv2d_transpose(
+        conv3d = self.block_builder.emit(
+            relax.op.nn.conv3d(
                 x,
                 weight,
                 strides=module.stride,
                 padding=module.padding,
                 dilation=module.dilation,
                 groups=module.groups,
-                data_layout="NCHW",
-                kernel_layout="OIHW",
+                data_layout="NCDHW",
+                kernel_layout="OIDHW",
                 out_dtype="float32",
             )
         )
 
         if module.bias is None:
-            return conv2d_transpose
+            return conv3d
 
         bias = self.params[module.bias]
         assert len(self.shape_of(bias)) == 1
-        bias = relax.op.reshape(bias, (1, -1, 1, 1))
-
-        return self.block_builder.emit(relax.op.add(conv2d_transpose, bias))
-
-    def _conv2d(self, node: fx.node.Node) -> relax.Var:
-        x = self.env[node.args[0]]
-        module = self.named_modules[node.target]
-        weight = self.params[module.weight]
-        bias = None
-        if module.bias is not None:
-            bias = self.params[module.bias]
-
-        return self._conv2d_impl(
-            x,
-            weight,
-            bias=bias,
-            strides=module.stride,
-            padding=module.padding,
-            dilation=module.dilation,
-            groups=module.groups,
-        )
+        bias = relax.op.reshape(bias, (1, -1, 1, 1, 1))
 
-    def _conv2d_functional(self, node: fx.node.Node) -> relax.Var:
-        args = self.retrieve_args(node)
-        x = args[0]
-        weight = args[1]
-        bias = args[2] if len(args) > 2 else None
-        stride = args[3] if len(args) > 3 else 1
-        padding = args[4] if len(args) > 4 else 0
-        dilation = args[5] if len(args) > 5 else 1
-        groups = args[6] if len(args) > 6 else 1
-        return self._conv2d_impl(
-            x,
-            weight,
-            bias=bias,
-            strides=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-        )
+        return self.block_builder.emit(relax.op.add(conv3d, bias))
 
     def _max_pool2d(self, node: fx.node.Node) -> relax.Var:
         x = self.env[node.args[0]]