fix einsum squeeze (#486)

Author: joey12300

paddlenlp/ops/einsum.py

@@ -348,5 +348,6 @@ def _mul_sum(left, right, sum_dims):
     squeeze_dims = [
         i for i in range(len(result.shape) - 1, num_output_dims - 1, -1)
     ]
-    result = paddle.squeeze(result, squeeze_dims)
+    if len(squeeze_dims) != 0:
+        result = paddle.squeeze(result, squeeze_dims)
     return result

tests/utils/test_ops.py

@@ -14,7 +14,7 @@
 import numpy as np
 import unittest
 import paddle
-import paddlenlp.utils.ops as ops
+import paddlenlp.ops as ops
 from common_test import CommonTest
 
 EINSUM_TEST_SAMPLE = {
@@ -29,6 +29,8 @@
     "G": np.random.rand(4, 2, 5),
     "H": np.random.rand(3, 2, 4),
     "I": np.random.rand(2, 2),
+    "J": np.random.rand(1, 3, 5),
+    "K": np.random.rand(1, 2, 3, 4),
 }
 
 
@@ -171,5 +173,10 @@ def setUp(self):
         self.sample = {"paradigm": "ijkl, lmn->ijn", "data": ["F", "H"]}
 
 
+class TestEinsumBatch1(TestEinsum):
+    def setUp(self):
+        self.sample = {"paradigm": "blq,bhlk->bhlqk", "data": ["J", "K"]}
+
+
 if __name__ == "__main__":
     unittest.main()