[Relax][PyTorch] Add stack.default and sum.default to exported programs translator (#17814)

* stack correct

* sum correct in side script

* all pass

Author: hugolatendresse

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

@@ -381,6 +381,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
@@ -409,6 +410,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,

tests/python/relax/test_from_exported_to_cuda.py

@@ -512,5 +512,33 @@ 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)
+
+
+@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()