[PIR]Migrate prod into pir (PaddlePaddle#58105)

danleifeng · Oct 25, 2023 · 8d0c452 · 8d0c452
1 parent 1c1f1d2
commit 8d0c452
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 32 deletions.
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -130,7 +130,7 @@ def _get_reduce_axis(axis, x):
 
 
 def _get_reduce_axis_with_tensor(axis, x):
-    if isinstance(axis, Variable):
+    if isinstance(axis, (Variable, paddle.pir.OpResult)):
         if axis.shape[0] == len(x.shape):
             reduce_all = True
         else:
@@ -4538,7 +4538,7 @@ def prod(x, axis=None, keepdim=False, dtype=None, name=None):
             x = cast(x, dtype)
 
     reduce_all, axis = _get_reduce_axis_with_tensor(axis, x)
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.prod(x, axis, keepdim, reduce_all)
     else:
         helper = LayerHelper('reduce_prod', **locals())

diff --git a/test/legacy_test/test_prod_op.py b/test/legacy_test/test_prod_op.py
@@ -18,6 +18,7 @@
 from test_sum_op import TestReduceOPTensorAxisBase
 
 import paddle
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestProdOp(unittest.TestCase):
@@ -70,33 +71,35 @@ def run_imperative(self):
             dy_result.numpy(), expected_result, rtol=1e-05
         )
 
+    @test_with_pir_api
     def run_static(self, use_gpu=False):
-        input = paddle.static.data(
-            name='input', shape=[10, 10, 5], dtype='float32'
-        )
-        result0 = paddle.prod(input)
-        result1 = paddle.prod(input, axis=1)
-        result2 = paddle.prod(input, axis=-1)
-        result3 = paddle.prod(input, axis=[0, 1])
-        result4 = paddle.prod(input, axis=1, keepdim=True)
-        result5 = paddle.prod(input, axis=1, dtype='int64')
-        result6 = paddle.prod(input, axis=1, keepdim=True, dtype='int64')
-
-        place = paddle.CUDAPlace(0) if use_gpu else paddle.CPUPlace()
-        exe = paddle.static.Executor(place)
-        exe.run(paddle.static.default_startup_program())
-        static_result = exe.run(
-            feed={"input": self.input},
-            fetch_list=[
-                result0,
-                result1,
-                result2,
-                result3,
-                result4,
-                result5,
-                result6,
-            ],
-        )
+        with paddle.static.program_guard(paddle.static.Program()):
+            input = paddle.static.data(
+                name='input', shape=[10, 10, 5], dtype='float32'
+            )
+            result0 = paddle.prod(input)
+            result1 = paddle.prod(input, axis=1)
+            result2 = paddle.prod(input, axis=-1)
+            result3 = paddle.prod(input, axis=[0, 1])
+            result4 = paddle.prod(input, axis=1, keepdim=True)
+            result5 = paddle.prod(input, axis=1, dtype='int64')
+            result6 = paddle.prod(input, axis=1, keepdim=True, dtype='int64')
+
+            place = paddle.CUDAPlace(0) if use_gpu else paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            exe.run(paddle.static.default_startup_program())
+            static_result = exe.run(
+                feed={"input": self.input},
+                fetch_list=[
+                    result0,
+                    result1,
+                    result2,
+                    result3,
+                    result4,
+                    result5,
+                    result6,
+                ],
+            )
 
         expected_result = np.prod(self.input)
         np.testing.assert_allclose(
@@ -134,8 +137,7 @@ def test_cpu(self):
         self.run_imperative()
         paddle.enable_static()
 
-        with paddle.static.program_guard(paddle.static.Program()):
-            self.run_static()
+        self.run_static()
 
     def test_gpu(self):
         if not paddle.base.core.is_compiled_with_cuda():
@@ -145,8 +147,7 @@ def test_gpu(self):
         self.run_imperative()
         paddle.enable_static()
 
-        with paddle.static.program_guard(paddle.static.Program()):
-            self.run_static(use_gpu=True)
+        self.run_static(use_gpu=True)
 
 
 class TestProdOpError(unittest.TestCase):