Initial Commit

Author: Netra-Srinivasan-McW

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

@@ -2107,7 +2107,26 @@ def _broadcast_to(self, node: fx.Node) -> relax.Var:
     def _cat(self, node: fx.Node) -> relax.Var:
         args = self.retrieve_args(node)
         axis = args[1] if len(node.args) > 1 else node.kwargs.get("dim", 0)
-        return self.block_builder.emit(relax.op.concat(args[0], axis=axis))
+        tensors = args[0]
+
+        filtered = []
+        for t in tensors:
+            shape = self.shape_of(t)
+            if shape is not None and len(shape) == 1:
+                dim0 = self._static_dim(shape[0])
+                if dim0 == 0:
+                    continue
+            filtered.append(t)
+
+        if not filtered:
+            # Passes args[0] to relax.op.concat
+            filtered = tensors 
+
+        if len(filtered) == 1:
+            # Returns the only tensor that is not a 1D empty tensor
+            return filtered[0] 
+
+        return self.block_builder.emit(relax.op.concat(filtered, axis=axis))
 
     def _chunk(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]

tests/python/relax/test_frontend_from_exported_program.py

@@ -4984,6 +4984,10 @@ class Cat1(Module):
         def forward(self, x, y):
             return torch.cat((x, y), dim=1)
 
+    class Cat2(Module):
+        def forward(self, x, y):
+            return torch.cat((x, y), dim=0)
+
     @I.ir_module
     class Expected1:
         @R.function
@@ -5010,9 +5014,24 @@ def main(
                 R.output(gv)
             return gv
 
+    # Expected IR: Concatenate a 1D empty tensor with any tensor
+    @I.ir_module
+    class Expected3:
+        @R.function
+        def main(
+            inp_0: R.Tensor((0,), dtype="float32"),
+            inp_1: R.Tensor((2, 3), dtype="float32"),
+        ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")):
+            with R.dataflow():
+                gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (inp_1,)
+                R.output(gv)
+            return gv
+
     example_args = (torch.randn(2, 3, dtype=torch.float32), torch.randn(2, 3, dtype=torch.float32))
+    example_args_empty = (torch.tensor([], dtype=torch.float32), torch.randn(2, 3, dtype=torch.float32))
     verify_model(Cat0(), example_args, {}, Expected1)
     verify_model(Cat1(), example_args, {}, Expected2)
+    verify_model(Cat2(), example_args_empty, {}, Expected3)
 
 
 def test_cumsum():

tests/python/relax/test_frontend_from_fx.py

@@ -5074,10 +5074,24 @@ def main(
                 R.output(gv)
             return gv
 
+    # Expected IR: Concatenate a 1D empty tensor with any tensor
+    @I.ir_module
+    class Expected3:
+        @R.function
+        def main(
+            inp_0: R.Tensor((0,), dtype="float32"),
+            inp_1: R.Tensor((2, 3), dtype="float32"),
+        ) -> R.Tensor((2, 3), dtype="float32"):
+            with R.dataflow():
+                gv: R.Tensor((2, 3), dtype="float32") = inp_1
+                R.output(gv)
+            return gv
+
     verify_model(Cat0(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)
     verify_model(Cat1(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
     verify_model(Cat2(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
     verify_model(Cat3(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)
+    verify_model(Cat3(), [([0], "float32"), ([2, 3], "float32")], {}, Expected3)
 
 
 def test_max():