From c839bcc7e62a4ee95e2067a4d11080e98bf74915 Mon Sep 17 00:00:00 2001 From: Hugo Latendresse Date: Mon, 7 Apr 2025 12:12:01 -0400 Subject: [PATCH 1/3] stack correct --- .../torch/exported_program_translator.py | 1 + tests/python/relax/test_from_exported_to_cuda.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index cc9217c9f5f8..7cf0901f30bc 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -389,6 +389,7 @@ def create_convert_map( "split_with_sizes.default": self._split, "squeeze.default": self._squeeze, "squeeze.dim": self._squeeze, + "stack.default": self._stack, "take.default": self._take, "tile.default": self._tile, "topk.default": self._topk, diff --git a/tests/python/relax/test_from_exported_to_cuda.py b/tests/python/relax/test_from_exported_to_cuda.py index 56ee527caf09..1b9c77229863 100644 --- a/tests/python/relax/test_from_exported_to_cuda.py +++ b/tests/python/relax/test_from_exported_to_cuda.py @@ -512,5 +512,21 @@ def forward(self, x): assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) +@tvm.testing.parametrize_targets("cuda") +def test_stack(target, dev): + class StackModel(nn.Module): + def forward(self, x): + val1 = x[1,4] + val2 = x[3,2] + val3 = x[5,6] + z = torch.stack([val1, val2, val3]) + return z + + torch_module = StackModel().eval() + raw_data = np.random.rand(10,10,10).astype("float32") + assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) + + + if __name__ == "__main__": tvm.testing.main() From 52689bc185b540c1fcf42ce8a8e47f358ba9091c Mon Sep 17 00:00:00 2001 From: Hugo Latendresse Date: Mon, 7 Apr 2025 12:14:41 -0400 Subject: [PATCH 2/3] sum correct in side script --- python/tvm/relax/frontend/torch/exported_program_translator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 7cf0901f30bc..37e8e352c6a7 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -361,6 +361,7 @@ def create_convert_map( "mean.dim": self._mean, "prod.default": self._prod, "std.correction": self._std, + "sum.default": self._sum, "sum.dim_IntList": self._sum, "var.correction": self._var, # search From 8d0e5bbff7b93502670afceaa8241b1e2bab75a4 Mon Sep 17 00:00:00 2001 From: Hugo Latendresse Date: Mon, 7 Apr 2025 12:17:20 -0400 Subject: [PATCH 3/3] all pass --- .../relax/test_from_exported_to_cuda.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/tests/python/relax/test_from_exported_to_cuda.py b/tests/python/relax/test_from_exported_to_cuda.py index 1b9c77229863..8405f48576d8 100644 --- a/tests/python/relax/test_from_exported_to_cuda.py +++ b/tests/python/relax/test_from_exported_to_cuda.py @@ -516,17 +516,29 @@ def forward(self, x): def test_stack(target, dev): class StackModel(nn.Module): def forward(self, x): - val1 = x[1,4] - val2 = x[3,2] - val3 = x[5,6] + val1 = x[1, 4] + val2 = x[3, 2] + val3 = x[5, 6] z = torch.stack([val1, val2, val3]) return z - + torch_module = StackModel().eval() - raw_data = np.random.rand(10,10,10).astype("float32") + raw_data = np.random.rand(10, 10, 10).astype("float32") assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) +@tvm.testing.parametrize_targets("cuda") +def test_sum(target, dev): + class SumModel(nn.Module): + def forward(self, x): + new_vec = x[1, 4] + return new_vec.sum() + + torch_module = SumModel().eval() + + raw_data = np.random.rand(10, 10, 10).astype("float32") + assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, target, dev) + if __name__ == "__main__": tvm.testing.main()