From d982c7896b245afa048ac2b3bf57ffbfc4090f2b Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Sun, 2 Nov 2025 00:27:32 -0400 Subject: [PATCH 1/3] f1 --- .../test_frontend_from_exported_program.py | 113 ++++++++++-------- 1 file changed, 65 insertions(+), 48 deletions(-) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 019d64955857..13781ca2f372 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5222,17 +5222,17 @@ def forward(self, data): class Expected: @R.function def main( - inp_0: R.Tensor((5,), dtype="float32"), + data: R.Tensor((5,), dtype="float32"), ) -> R.Tuple(R.Tensor((5,), dtype="float32")): with R.dataflow(): - lv: R.Tensor((5,), dtype="float32") = R.zeros_like(inp_0, dtype="void") + lv: R.Tensor((5,), dtype="float32") = R.zeros(R.shape([5]), dtype="float32") gv: R.Tuple(R.Tensor((5,), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.randn(5, dtype=torch.float32),) - verify_model(EmptyLike(), example_args, {}, Expected) + verify_model(EmptyLike(), example_args, {}, Expected, run_ep_decomposition=True) def test_one_hot(): @@ -5244,19 +5244,22 @@ def forward(self, indices): class Expected: @R.function def main( - inp_0: R.Tensor((5,), dtype="int64"), + indices: R.Tensor((5,), dtype="int64"), ) -> R.Tuple(R.Tensor((5, 10), dtype="int64")): with R.dataflow(): - lv: R.Tensor((5, 10), dtype="int64") = R.one_hot( - inp_0, R.prim_value(1), R.prim_value(0), depth=10, axis=-1 + lv: R.Tensor((10,), dtype="int64") = R.arange( + R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64" ) - gv: R.Tuple(R.Tensor((5, 10), dtype="int64")) = (lv,) + lv1: R.Tensor((5, 1), dtype="int64") = R.expand_dims(indices, axis=[-1]) + lv2: R.Tensor((5, 10), dtype="bool") = R.equal(lv1, lv) + lv3: R.Tensor((5, 10), dtype="int64") = R.astype(lv2, dtype="int64") + gv: R.Tuple(R.Tensor((5, 10), dtype="int64")) = (lv3,) R.output(gv) return gv example_args = (torch.randint(0, 10, (5,), dtype=torch.int64),) - verify_model(OneHot(), example_args, {}, Expected) + verify_model(OneHot(), example_args, {}, Expected, run_ep_decomposition=True) def test_ones_like(): @@ -5271,14 +5274,16 @@ def main( input: R.Tensor((128, 128), dtype="float32") ) -> R.Tuple(R.Tensor((128, 128), dtype="float32")): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.ones_like(input, dtype="void") + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(1, "int32"), dtype="void" + ) gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.rand(128, 128, dtype=torch.float32),) - verify_model(OnesLike(), example_args, {}, Expected) + verify_model(OnesLike(), example_args, {}, Expected, run_ep_decomposition=True) def test_zero_inplace(): @@ -5291,16 +5296,23 @@ class Expected: @R.function def main( input: R.Tensor((128, 128), dtype="float32") - ) -> R.Tuple(R.Tensor((128, 128), dtype="float32")): + ) -> R.Tuple(R.Tensor((128, 128), dtype="float32"), R.Tensor((128, 128), dtype="float32")): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(input, dtype="void") - gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(0, "int32"), dtype="void" + ) + gv: R.Tuple( + R.Tensor((128, 128), dtype="float32"), R.Tensor((128, 128), dtype="float32") + ) = ( + lv, + lv, + ) R.output(gv) return gv example_args = (torch.rand(128, 128, dtype=torch.float32),) - verify_model(ZeroInplace(), example_args, {}, Expected) + verify_model(ZeroInplace(), example_args, {}, Expected, run_ep_decomposition=True) def test_zeros(): @@ -5315,14 +5327,16 @@ def main( input: R.Tensor((128, 128), dtype="float32") ) -> R.Tuple(R.Tensor((5, 2), dtype="float32")): with R.dataflow(): - lv: R.Tensor((5, 2), dtype="float32") = R.zeros(R.shape([5, 2]), dtype="float32") + lv: R.Tensor((5, 2), dtype="float32") = R.full( + R.shape([5, 2]), R.const(0.0, "float32"), dtype="float32" + ) gv: R.Tuple(R.Tensor((5, 2), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.rand(128, 128, dtype=torch.float32),) - verify_model(Zeros(), example_args, {}, Expected) + verify_model(Zeros(), example_args, {}, Expected, run_ep_decomposition=True) def test_zeros_like(): @@ -5337,13 +5351,15 @@ def main( input: R.Tensor((128, 128), dtype="float32") ) -> R.Tuple(R.Tensor((128, 128), dtype="float32")): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(input, dtype="void") + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(0, "int32"), dtype="void" + ) gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.rand(128, 128, dtype=torch.float32),) - verify_model(ZerosLike(), example_args, {}, Expected) + verify_model(ZerosLike(), example_args, {}, Expected, run_ep_decomposition=True) def test_type_as(): @@ -5369,7 +5385,7 @@ def main( torch.rand(128, 128, dtype=torch.float16), ) - verify_model(TypeAs(), example_args, {}, Expected) + verify_model(TypeAs(), example_args, {}, Expected, run_ep_decomposition=True) def test_select(): @@ -5391,7 +5407,7 @@ def main( example_args = (torch.randn(2, 3, dtype=torch.float32),) - verify_model(Select(), example_args, {}, Expected) + verify_model(Select(), example_args, {}, Expected, run_ep_decomposition=True) def test_unflatten(): @@ -5417,8 +5433,8 @@ def main( example_args = (torch.randn(2, 15, 7, dtype=torch.float32),) - verify_model(Unflatten(), example_args, {}, Expected) - verify_model(Unflatten1(), example_args, {}, Expected) + verify_model(Unflatten(), example_args, {}, Expected, run_ep_decomposition=True) + verify_model(Unflatten1(), example_args, {}, Expected, run_ep_decomposition=True) def test_gather(): @@ -5495,10 +5511,10 @@ def main( torch.randint(0, 3, (2, 3), dtype=torch.int64), ) - verify_model(Gather0(), example_args, {}, Expected0) - verify_model(Gather1(), example_args, {}, Expected1) - verify_model(Gather2(), example_args, {}, Expected2) - verify_model(Gather3(), example_args, {}, Expected3) + verify_model(Gather0(), example_args, {}, Expected0, run_ep_decomposition=True) + verify_model(Gather1(), example_args, {}, Expected1, run_ep_decomposition=True) + verify_model(Gather2(), example_args, {}, Expected2, run_ep_decomposition=True) + verify_model(Gather3(), example_args, {}, Expected3, run_ep_decomposition=True) def test_index_put(): @@ -5669,11 +5685,11 @@ def main( return gv # Run verification for each case - verify_model(IndexPut1D(), example_args_1d, {}, Expected1D) - verify_model(IndexPut2D(), example_args_2d, {}, Expected2D) - verify_model(IndexPut3D(), example_args_3d, {}, Expected3D) - verify_model(IndexPut4D(), example_args_4d, {}, Expected4D) - verify_model(IndexPut5D(), example_args_5d, {}, Expected5D) + verify_model(IndexPut1D(), example_args_1d, {}, Expected1D, run_ep_decomposition=True) + verify_model(IndexPut2D(), example_args_2d, {}, Expected2D, run_ep_decomposition=True) + verify_model(IndexPut3D(), example_args_3d, {}, Expected3D, run_ep_decomposition=True) + verify_model(IndexPut4D(), example_args_4d, {}, Expected4D, run_ep_decomposition=True) + verify_model(IndexPut5D(), example_args_5d, {}, Expected5D, run_ep_decomposition=True) def test_flip(): @@ -5711,8 +5727,8 @@ def main( example_args = (torch.randn(2, 2, dtype=torch.float32),) - verify_model(Flip0(), example_args, {}, Expected0) - verify_model(Flip1(), example_args, {}, Expected1) + verify_model(Flip0(), example_args, {}, Expected0, run_ep_decomposition=True) + verify_model(Flip1(), example_args, {}, Expected1, run_ep_decomposition=True) def test_take(): @@ -5724,12 +5740,12 @@ def forward(self, data, indices): class Expected: @R.function def main( - inp_0: R.Tensor((5,), dtype="float32"), - inp_1: R.Tensor((3,), dtype="int64"), + data: R.Tensor((5,), dtype="float32"), + indices: R.Tensor((3,), dtype="int64"), ) -> R.Tuple(R.Tensor((3,), dtype="float32")): with R.dataflow(): - lv: R.Tensor((3,), dtype="int32") = R.astype(inp_1, dtype="int32") - lv1: R.Tensor((3,), dtype="float32") = R.take(inp_0, lv, axis=None) + lv: R.Tensor((5,), dtype="float32") = R.reshape(data, R.shape([5])) + lv1: R.Tensor((3,), dtype="float32") = R.index_tensor(lv, (indices,)) gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv1,) R.output(gv) return gv @@ -5739,7 +5755,7 @@ def main( torch.randint(0, 5, (3,), dtype=torch.int64), ) - verify_model(Take(), example_args, {}, Expected) + verify_model(Take(), example_args, {}, Expected, run_ep_decomposition=True) def test_std(): @@ -5751,16 +5767,17 @@ def forward(self, x): class Expected: @R.function def main( - inp_0: R.Tensor((5, 3), dtype="float32"), + x: R.Tensor((5, 3), dtype="float32"), ) -> R.Tuple(R.Tensor((), dtype="float32")): with R.dataflow(): - lv: R.Tensor((), dtype="float32") = R.std(inp_0, axis=None, keepdims=False) - gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,) + lv: R.Tensor((), dtype="float32") = R.variance(x, axis=None, keepdims=False) + lv1: R.Tensor((), dtype="float32") = R.sqrt(lv) + gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv1,) R.output(gv) return gv example_args = (torch.randn(5, 3, dtype=torch.float32),) - verify_model(Std(), example_args, {}, Expected) + verify_model(Std(), example_args, {}, Expected, run_ep_decomposition=True) def test_var(): @@ -5772,16 +5789,16 @@ def forward(self, x): class Expected: @R.function def main( - inp_0: R.Tensor((5, 3), dtype="float32"), + x: R.Tensor((5, 3), dtype="float32"), ) -> R.Tuple(R.Tensor((), dtype="float32")): with R.dataflow(): - lv: R.Tensor((), dtype="float32") = R.variance(inp_0, axis=None, keepdims=False) + lv: R.Tensor((), dtype="float32") = R.variance(x, axis=None, keepdims=False) gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.randn(5, 3, dtype=torch.float32),) - verify_model(Var(), example_args, {}, Expected) + verify_model(Var(), example_args, {}, Expected, run_ep_decomposition=True) def test_prod(): @@ -5793,16 +5810,16 @@ def forward(self, x): class Expected: @R.function def main( - inp_0: R.Tensor((5, 3), dtype="float32"), + x: R.Tensor((5, 3), dtype="float32"), ) -> R.Tuple(R.Tensor((), dtype="float32")): with R.dataflow(): - lv: R.Tensor((), dtype="float32") = R.prod(inp_0, axis=None, keepdims=False) + lv: R.Tensor((), dtype="float32") = R.prod(x, axis=None, keepdims=False) gv: R.Tuple(R.Tensor((), dtype="float32")) = (lv,) R.output(gv) return gv example_args = (torch.randn(5, 3, dtype=torch.float32),) - verify_model(Prod(), example_args, {}, Expected) + verify_model(Prod(), example_args, {}, Expected, run_ep_decomposition=True) def test_cumprod(): From 7b98cdecefe8952ea4eab95283d497c931144962 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Sun, 2 Nov 2025 13:32:49 -0500 Subject: [PATCH 2/3] f2 --- python/tvm/relax/frontend/torch/exported_program_translator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 48ae002c05c0..eda20ef502a0 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1003,6 +1003,7 @@ def create_convert_map( "flip.default": self._flip, "gather.default": self._gather, "index.Tensor": self._index_tensor, + "index_put.default": self._index_put, "index_put_.default": self._index_put, "meshgrid.indexing": self._meshgrid, "meshgrid.default": self._meshgrid, @@ -1041,6 +1042,7 @@ def create_convert_map( "contiguous.default": lambda node: self.env[node.args[0]], # no-op "clone.default": lambda node: self.env[node.args[0]], "bernoulli.p": lambda node: self.env[node.args[0]], # Dropout: just return input + "_assert_tensor_metadata.default": lambda node: self.env[node.args[0]], # metadata assertion: no-op "empty.memory_format": self._empty, "empty_permuted.default": self._empty, # Similar to empty with permuted layout "empty_like.default": self._empty_like, From b2e3ff90ab8ada9c6812795bd92ef1f8ca4b2189 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Sun, 2 Nov 2025 14:18:43 -0500 Subject: [PATCH 3/3] f4 --- .../torch/exported_program_translator.py | 4 +- .../test_frontend_from_exported_program.py | 41 ++++++++++++++----- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index eda20ef502a0..3be255a29a65 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1042,7 +1042,9 @@ def create_convert_map( "contiguous.default": lambda node: self.env[node.args[0]], # no-op "clone.default": lambda node: self.env[node.args[0]], "bernoulli.p": lambda node: self.env[node.args[0]], # Dropout: just return input - "_assert_tensor_metadata.default": lambda node: self.env[node.args[0]], # metadata assertion: no-op + "_assert_tensor_metadata.default": lambda node: self.env[ + node.args[0] + ], # metadata assertion: no-op "empty.memory_format": self._empty, "empty_permuted.default": self._empty, # Similar to empty with permuted layout "empty_like.default": self._empty_like, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 13781ca2f372..9f63743faa29 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5537,12 +5537,15 @@ def main( data: R.Tensor((64,), dtype="float32"), indices_0: R.Tensor((128,), dtype="int64"), values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((64,), dtype="float32")): + ) -> R.Tuple(R.Tensor((64,), dtype="float32"), R.Tensor((64,), dtype="float32")): with R.dataflow(): lv: R.Tensor((64,), dtype="float32") = R.index_put( data, R.tuple(indices_0), values, accumulate=False ) - gv: R.Tuple(R.Tensor((64,), dtype="float32")) = (lv,) + gv: R.Tuple(R.Tensor((64,), dtype="float32"), R.Tensor((64,), dtype="float32")) = ( + lv, + lv, + ) R.output(gv) return gv @@ -5567,12 +5570,14 @@ def main( indices_0: R.Tensor((128,), dtype="int64"), indices_1: R.Tensor((128,), dtype="int64"), values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((32, 64), dtype="float32")): + ) -> R.Tuple(R.Tensor((32, 64), dtype="float32"), R.Tensor((32, 64), dtype="float32")): with R.dataflow(): lv: R.Tensor((32, 64), dtype="float32") = R.index_put( data, R.tuple(indices_0, indices_1), values, accumulate=False ) - gv: R.Tuple(R.Tensor((32, 64), dtype="float32")) = (lv,) + gv: R.Tuple( + R.Tensor((32, 64), dtype="float32"), R.Tensor((32, 64), dtype="float32") + ) = (lv, lv) R.output(gv) return gv @@ -5599,12 +5604,16 @@ def main( indices_1: R.Tensor((128,), dtype="int64"), indices_2: R.Tensor((128,), dtype="int64"), values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((16, 32, 64), dtype="float32")): + ) -> R.Tuple( + R.Tensor((16, 32, 64), dtype="float32"), R.Tensor((16, 32, 64), dtype="float32") + ): with R.dataflow(): lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put( data, R.tuple(indices_0, indices_1, indices_2), values, accumulate=False ) - gv: R.Tuple(R.Tensor((16, 32, 64), dtype="float32")) = (lv,) + gv: R.Tuple( + R.Tensor((16, 32, 64), dtype="float32"), R.Tensor((16, 32, 64), dtype="float32") + ) = (lv, lv) R.output(gv) return gv @@ -5633,7 +5642,10 @@ def main( indices_2: R.Tensor((128,), dtype="int64"), indices_3: R.Tensor((128,), dtype="int64"), values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")): + ) -> R.Tuple( + R.Tensor((8, 16, 32, 64), dtype="float32"), + R.Tensor((8, 16, 32, 64), dtype="float32"), + ): with R.dataflow(): lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put( data, @@ -5641,7 +5653,10 @@ def main( values, accumulate=False, ) - gv: R.Tuple(R.Tensor((8, 16, 32, 64), dtype="float32")) = (lv,) + gv: R.Tuple( + R.Tensor((8, 16, 32, 64), dtype="float32"), + R.Tensor((8, 16, 32, 64), dtype="float32"), + ) = (lv, lv) R.output(gv) return gv @@ -5672,7 +5687,10 @@ def main( indices_3: R.Tensor((128,), dtype="int64"), indices_4: R.Tensor((128,), dtype="int64"), values: R.Tensor((128,), dtype="float32"), - ) -> R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")): + ) -> R.Tuple( + R.Tensor((4, 8, 16, 32, 64), dtype="float32"), + R.Tensor((4, 8, 16, 32, 64), dtype="float32"), + ): with R.dataflow(): lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put( data, @@ -5680,7 +5698,10 @@ def main( values, accumulate=False, ) - gv: R.Tuple(R.Tensor((4, 8, 16, 32, 64), dtype="float32")) = (lv,) + gv: R.Tuple( + R.Tensor((4, 8, 16, 32, 64), dtype="float32"), + R.Tensor((4, 8, 16, 32, 64), dtype="float32"), + ) = (lv, lv) R.output(gv) return gv