add unit test

Author: Archermmt

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

@@ -604,7 +604,7 @@ def _split(self, node: fx.node.Node) -> relax.Var:
             dim = 0
         if isinstance(split_size, (list, tuple)):
             n_section = []
-            for s in split_size:
+            for s in split_size[:-1]:
                 cum_sum = 0 if not n_section else n_section[-1]
                 n_section.append(s + cum_sum)
         else:

tests/python/contrib/test_msc/test_graph_build.py

@@ -1345,11 +1345,15 @@ def forward(self, x_1, x_2, x_3):
 def test_split():
     """test graph builder for split"""
 
-    class Split(Module):
+    class Split1(Module):
         def forward(self, data):
             return torch.split(data, 1, dim=1)
 
-    expected = {
+    class Split2(Module):
+        def forward(self, data):
+            return torch.split(data, [1, 2], dim=1)
+
+    expected1 = {
         "inputs": [
             {"name": "inp_0", "shape": [1, 3, 10, 10], "dtype": "float32", "layout": "ABCD"}
         ],
@@ -1361,8 +1365,20 @@ def forward(self, data):
         "nodes": {"total": 2, "input": 1, "split": 1},
     }
 
+    expected2 = {
+        "inputs": [
+            {"name": "inp_0", "shape": [1, 3, 10, 10], "dtype": "float32", "layout": "ABCD"}
+        ],
+        "outputs": [
+            {"name": "split_0", "shape": [1, 1, 10, 10], "dtype": "float32", "layout": "ABCD"},
+            {"name": "split_1", "shape": [1, 2, 10, 10], "dtype": "float32", "layout": "ABCD"},
+        ],
+        "nodes": {"total": 2, "input": 1, "split": 1},
+    }
+
     input_info = [([1, 3, 10, 10], "float32")]
-    verify_model(Split(), input_info, expected)
+    verify_model(Split1(), input_info, expected1)
+    verify_model(Split2(), input_info, expected2)
 
 
 def test_unbind():
@@ -1570,10 +1586,14 @@ def forward(self, x):
 def test_expand():
     """test graph builder for expand"""
 
-    class Expand(Module):
+    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)
+
     expected = {
         "inputs": [{"name": "inp_0", "shape": [1, 2, 3, 4], "dtype": "float32", "layout": ""}],
         "outputs": [
@@ -1583,7 +1603,8 @@ def forward(self, x):
     }
 
     input_info = [([1, 2, 3, 4], "float32")]
-    verify_model(Expand(), input_info, expected)
+    verify_model(Expand1(), input_info, expected)
+    verify_model(Expand2(), input_info, expected)
 
 
 def test_reduce():

tests/python/contrib/test_msc/test_translate_relax.py

@@ -67,7 +67,12 @@ def _run_relax(relax_mod):
 
     orig_output = _run_relax(orig_mod)
     rt_output = _run_relax(rt_mod)
-    tvm.testing.assert_allclose(orig_output, rt_output)
+    if not isinstance(orig_output, (list, tuple)):
+        orig_output = [orig_output]
+    if not isinstance(rt_output, (list, tuple)):
+        rt_output = [rt_output]
+    for o_out, r_out in zip(orig_output, rt_output):
+        tvm.testing.assert_allclose(o_out, r_out)
 
 
 def test_conv1d():
@@ -750,12 +755,17 @@ def forward(self, x_1, x_2, x_3):
 def test_split():
     """test relax translator for split"""
 
-    class Split(Module):
+    class Split1(Module):
         def forward(self, data):
             return torch.split(data, 1, dim=1)
 
+    class Split2(Module):
+        def forward(self, data):
+            return torch.split(data, [1, 2], dim=1)
+
     input_info = [([1, 3, 10, 10], "float32")]
-    _verify_model(Split(), input_info)
+    _verify_model(Split1(), input_info)
+    _verify_model(Split2(), input_info)
 
 
 def test_unbind():
@@ -890,12 +900,17 @@ def forward(self, x):
 def test_expand():
     """test relax translator for expand"""
 
-    class Expand(Module):
+    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)
+
     input_info = [([1, 2, 3, 4], "float32")]
-    _verify_model(Expand(), input_info)
+    _verify_model(Expand1(), input_info)
+    _verify_model(Expand2(), input_info)
 
 
 def test_reduce():

tests/python/contrib/test_msc/test_translate_relay.py

@@ -731,12 +731,17 @@ def forward(self, x_1, x_2, x_3):
 def test_split():
     """test relay to relax for split"""
 
-    class Split(Module):
+    class Split1(Module):
         def forward(self, data):
             return torch.split(data, 1, dim=1)
 
+    class Split2(Module):
+        def forward(self, data):
+            return torch.split(data, [1, 2], dim=1)
+
     input_info = [([1, 3, 10, 10], "float32")]
-    verify_model(Split(), input_info, build_target="llvm")
+    verify_model(Split1(), input_info, build_target="llvm")
+    verify_model(Split2(), input_info, build_target="llvm")
 
 
 def test_unbind():
@@ -875,12 +880,17 @@ def forward(self, x):
 def test_expand():
     """test relay to relax for expand"""
 
-    class Expand(Module):
+    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)
+
     input_info = [([1, 2, 3, 4], "float32")]
-    verify_model(Expand(), input_info, build_target="llvm")
+    verify_model(Expand1(), input_info, build_target="llvm")
+    verify_model(Expand2(), input_info, build_target="llvm")
 
 
 def test_reduce():

tests/python/contrib/test_msc/test_translate_tensorrt.py

@@ -673,12 +673,17 @@ def forward(self, x_1, x_2, x_3):
 def test_split():
     """test tensorrt translator for split"""
 
-    class Split(Module):
+    class Split1(Module):
         def forward(self, data):
             return torch.split(data, 1, dim=1)
 
+    class Split2(Module):
+        def forward(self, data):
+            return torch.split(data, [1, 2], dim=1)
+
     input_info = [([1, 3, 10, 10], "float32")]
-    verify_model(Split(), input_info)
+    verify_model(Split1(), input_info)
+    verify_model(Split2(), input_info)
 
 
 @requires_tensorrt
@@ -714,13 +719,19 @@ def forward(self, data):
 def test_expand():
     """test tensorrt translator for expand"""
 
-    class Expand(Module):
+    class Expand1(Module):
         def forward(self, x):
             x = x + 1.0
             return x.expand(4, 2, 3, 4)
 
+    class Expand2(Module):
+        def forward(self, x):
+            x = x + 1.0
+            return x.expand(4, -1, -1, 4)
+
     input_info = [([1, 2, 3, 4], "float32")]
-    verify_model(Expand(), input_info)
+    verify_model(Expand1(), input_info)
+    verify_model(Expand2(), input_info)
 
 
 @requires_tensorrt

tests/python/contrib/test_msc/test_translate_torch.py

@@ -728,13 +728,18 @@ def forward(self, x_1, x_2, x_3):
 def test_split():
     """test torch translator for split"""
 
-    class Split(Module):
+    class Split1(Module):
         def forward(self, data):
             return torch.split(data, 1, dim=1)
 
+    class Split2(Module):
+        def forward(self, data):
+            return torch.split(data, [1, 2], dim=1)
+
     input_info = [([1, 3, 10, 10], "float32")]
     for via_relax in [True, False]:
-        verify_model(Split(), input_info, via_relax)
+        verify_model(Split1(), input_info, via_relax)
+        verify_model(Split2(), input_info, via_relax)
 
 
 def test_unbind():
@@ -852,13 +857,18 @@ def forward(self, x):
 def test_expand():
     """test torch translator for expand"""
 
-    class Expand(Module):
+    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)
+
     input_info = [([1, 2, 3, 4], "float32")]
     for via_relax in [True, False]:
-        verify_model(Expand(), input_info, via_relax)
+        verify_model(Expand1(), input_info, via_relax)
+        verify_model(Expand2(), input_info, via_relax)
 
 
 def test_reduce():