【PIR API adaptor No.147】mm (PaddlePaddle#58935)

SecretXV · Nov 28, 2023 · 6ba9386 · 6ba9386
1 parent 646454b
commit 6ba9386
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -2128,7 +2128,7 @@ def mm(input, mat2, name=None):
 
 
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.matmul(input, mat2, False, False)
     else:
 

diff --git a/test/legacy_test/test_matmul_op.py b/test/legacy_test/test_matmul_op.py
@@ -19,6 +19,7 @@
 
 import paddle
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 def generate_compatible_shapes(dim_X, dim_Y, transpose_X, transpose_Y):
@@ -98,7 +99,10 @@ def test_check_output(self):
 
     def test_check_grad_normal(self):
         self.check_grad(
-            ['X', 'Y'], 'Out', max_relative_error=1e-3, check_cinn=True
+            ['X', 'Y'],
+            'Out',
+            max_relative_error=1e-3,
+            check_cinn=True,
         )
 
     def test_check_grad_ignore_x(self):
@@ -166,9 +170,10 @@ def generate_compatible_shapes_ndim(dim, transpose_X, transpose_Y):
 
 
 class API_TestMm(unittest.TestCase):
+    @test_with_pir_api
     def test_out(self):
         with paddle_static_guard():
-            with base.program_guard(base.Program()):
+            with paddle.base.program_guard(paddle.base.Program()):
                 x = paddle.static.data(name="x", shape=[2], dtype="float64")
                 y = paddle.static.data(name='y', shape=[2], dtype='float64')
                 result = paddle.mm(x, y)
@@ -220,7 +225,9 @@ def test_errors(self):
         with paddle_static_guard():
 
             def test_error1():
-                with base.program_guard(base.Program(), base.Program()):
+                with paddle.base.program_guard(
+                    paddle.base.Program(), paddle.base.Program()
+                ):
                     data1 = paddle.static.data(
                         name="data1", shape=[10, 2], dtype="float32"
                     )
@@ -232,7 +239,9 @@ def test_error1():
             self.assertRaises(ValueError, test_error1)
 
             def test_error2():
-                with base.program_guard(base.Program(), base.Program()):
+                with paddle.base.program_guard(
+                    paddle.base.Program(), paddle.base.Program()
+                ):
                     data1 = paddle.static.data(
                         name="data1", shape=[-1, 10, 2], dtype="float32"
                     )