diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 9bb3b77d73e9..70e0ef7e42ee 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2112,7 +2112,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]] diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 85fcf8b825ef..80078ece4842 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -4987,6 +4987,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 @@ -5013,9 +5017,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(): diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index d3c6315c132b..40bbc24d82a8 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -5110,10 +5110,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():