[PIR]Migrate einsum_v2 into pir (#58501)

0x45f · web-flow · commit 3f6163bc9c11 · 2023-10-31T14:54:08.000+08:00
diff --git a/python/paddle/tensor/einsum.py b/python/paddle/tensor/einsum.py
@@ -23,7 +23,7 @@
 from paddle import _C_ops
 
 from ..base.data_feeder import check_type, check_variable_and_dtype
-from ..base.framework import in_dygraph_mode
+from ..base.framework import in_dynamic_or_pir_mode
 from ..base.layer_helper import LayerHelper
 from .linalg import matmul, transpose
 from .manipulation import reshape, squeeze, unsqueeze
@@ -832,7 +832,7 @@ def gen_einsum_op(equation, *operands):
     EinsumOp Python Interface:
     """
 
-    if in_dygraph_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.einsum(operands, equation)[0]
     else:
         assert len(operands) <= 2, "Only support two operands in EinsumOp."
diff --git a/test/legacy_test/test_einsum_v2.py b/test/legacy_test/test_einsum_v2.py
@@ -19,6 +19,7 @@
 
 import paddle
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 os.environ['FLAGS_new_einsum'] = "1"
 
@@ -382,7 +383,7 @@ def check_output_equal(self, actual, expect, rtol=1.0e-5, atol=1.0e-8):
             rtol=rtol,
             atol=atol,
             err_msg=error_msg.format(
-                paddle.get_device(), expect, actual, self.__class__.__name__
+                self._get_place(False), expect, actual, self.__class__.__name__
             ),
         )
 
@@ -465,6 +466,7 @@ def test_sums(self):
         self.check_output("i,ij->", y, x)
         self.check_output("ij,i->", x, y)
 
+    @test_with_pir_api
     def test_static_graph(self):
         paddle.enable_static()
         base = paddle.base
@@ -523,11 +525,12 @@ def setUp(self):
     def tearDown(self):
         paddle.disable_static()
 
+    @test_with_pir_api
     def test_shape(self):
         A = paddle.static.data(name='x', shape=[-1])
         B = paddle.static.data(name='y', shape=[384])
         C = paddle.einsum('i,d->id', A, B)
-        self.assertEqual(C.shape, (-1, 384))
+        self.assertEqual(tuple(C.shape), (-1, 384))
 
 
 @unittest.skipIf(
