[feature] add different shape to unit tests (#8)

tests/test_llama.py

@@ -57,11 +57,19 @@ def __init__(self, hidden_size: int, intermediate_size: int):
         )
 
 
-class LlamaTestCase(unittest.TestCase):
-    def test_forward(self):
-        input = torch.zeros((1, 8, hidden_size))
-        output: torch.Tensor = moe_layer(input)
-        self.assertEqual(output.shape, (1, 8, hidden_size))
+class LlamaInputShapeTest(unittest.TestCase):
+    def test_forward_with_different_shape(self):
+        input_shapes = [
+            (2, 8, hidden_size),
+            (1, 16, hidden_size),
+            (4, 4, hidden_size)
+        ]
+
+        for shape in input_shapes:
+            with self.subTest(shape=shape):
+                input = torch.zeros(shape)
+                output: torch.Tensor = moe_layer(input)
+                self.assertEqual(output.shape, shape)
 
 
 if __name__ == "__main__":

tests/test_phi3.py

@@ -50,7 +50,7 @@ def __init__(self, hidden_size: int, intermediate_size: int):
         )
 
 
-class LlamaTestCase(unittest.TestCase):
+class Phi3TestCase(unittest.TestCase):
     def test_forward(self):
         input = torch.zeros((1, 8, hidden_size))
         output: torch.Tensor = moe_layer(input)