support expand

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

@@ -748,6 +748,17 @@ def _cumsum(self, node: fx.Node) -> relax.Var:
 
         return self.block_builder.emit(relax.op.cumsum(x, dim, dtype))
 
+    def _expand(self, node: fx.Node) -> relax.Var:
+        args = self.retrieve_args(node)
+        sizes = args[1:] if len(args) > 2 else args[1]
+        broadcast_shape, in_shape = [], self.shape_of(args[0])
+        for idx, i in enumerate(sizes):
+            if isinstance(i, int) and i == -1:
+                broadcast_shape.append(in_shape[idx])
+            else:
+                broadcast_shape.append(i)
+        return self.block_builder.emit(relax.op.broadcast_to(args[0], broadcast_shape))
+
     def _reshape(self, node: fx.Node) -> relax.Var:
         import torch  # type: ignore
 

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

@@ -252,6 +252,7 @@ def create_convert_map(
             "cat.default": self._cat,
             "concat.default": self._cat,
             "cumsum.default": self._cumsum,
+            "expand.default": self._expand,
             "view.default": self._reshape,
             # other
             "getitem": self._getitem,

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

@@ -386,17 +386,6 @@ def _chunk(self, node: fx.Node) -> relax.Var:
         dim = node.args[2] if len(node.args) > 2 else node.kwargs.get("dim", 0)
         return self.block_builder.emit(relax.op.split(x, chunks, dim))
 
-    def _expand(self, node: fx.Node) -> relax.Var:
-        args = self.retrieve_args(node)
-        sizes = args[1:] if len(args) > 2 else args[1]
-        broadcast_shape, in_shape = [], self.shape_of(args[0])
-        for idx, i in enumerate(sizes):
-            if isinstance(i, int) and i == -1:
-                broadcast_shape.append(in_shape[idx])
-            else:
-                broadcast_shape.append(i)
-        return self.block_builder.emit(relax.op.broadcast_to(args[0], broadcast_shape))
-
     def _flatten_impl(self, x, start_dim, end_dim) -> relax.Var:
         shape = self.shape_of(x)
         start_dim = start_dim if start_dim >= 0 else len(shape) + start_dim

tests/python/relax/test_frontend_from_exported_program.py

@@ -2806,6 +2806,33 @@ def main(
     verify_model(Cumsum(), example_args, {}, expected1)
 
 
+def test_expand():
+    class Expand1(Module):
+        def forward(self, x):
+            return x.expand(4, 2, 3, 4)
+
+    class Expand2(Module):
+        def forward(self, x):
+            return x.expand(4, -1, -1, 4)
+
+    @tvm.script.ir_module
+    class expected1:
+        @R.function
+        def main(
+            x: R.Tensor((1, 2, 3, 4), dtype="float32")
+        ) -> R.Tuple(R.Tensor((4, 2, 3, 4), dtype="float32")):
+            # block 0
+            with R.dataflow():
+                lv: R.Tensor((4, 2, 3, 4), dtype="float32") = R.broadcast_to(x, (4, 2, 3, 4))
+                gv: R.Tuple(R.Tensor((4, 2, 3, 4), dtype="float32")) = (lv,)
+                R.output(gv)
+            return gv
+
+    example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
+    verify_model(Expand1(), example_args, {}, expected1)
+    verify_model(Expand2(), example_args, {}, expected1)
+
+
 def test_view():
     class View(Module):
         def forward(self, x):