【HEU】[Paddle Tensor 第二期 复数类型支持] paddle.prod 支持复数类型 (#70546)

fangfangssj · web-flow · commit 8957f226a88b · 2025-01-03T16:08:28.000+08:00
* prod complex

* fix

* fix test

* fix
diff --git a/paddle/phi/kernels/cpu/prod_grad_kernel.cc b/paddle/phi/kernels/cpu/prod_grad_kernel.cc
@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/prod_grad_kernel.h"
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/common/complex.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/prod_grad_kernel_impl.h"
 
@@ -25,4 +26,6 @@ PD_REGISTER_KERNEL(prod_grad,
                    float,
                    double,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/prod_kernel.cc b/paddle/phi/kernels/cpu/prod_kernel.cc
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 #include "paddle/phi/kernels/prod_kernel.h"
-
 #include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/common/complex.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/reduce.h"
 #include "paddle/phi/kernels/funcs/reduce_functor.h"
@@ -36,5 +36,13 @@ void ProdKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    prod, CPU, ALL_LAYOUT, phi::ProdKernel, float, double, int, int64_t) {}
+PD_REGISTER_KERNEL(prod,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ProdKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/prod_grad_kernel.cu b/paddle/phi/kernels/gpu/prod_grad_kernel.cu
@@ -27,4 +27,6 @@ PD_REGISTER_KERNEL(prod_grad,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/kps/reduce_kernel.cu b/paddle/phi/kernels/kps/reduce_kernel.cu
@@ -415,5 +415,7 @@ PD_REGISTER_KERNEL(prod,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 #endif
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4788,7 +4788,16 @@ def prod(
         check_variable_and_dtype(
             x,
             'x/input',
-            ['float32', 'float64', 'int32', 'int64', "float16", "uint16"],
+            [
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                "float16",
+                "uint16",
+                "complex64",
+                "complex128",
+            ],
             'reduce_prod',
         )
         out = helper.create_variable_for_type_inference(
diff --git a/test/legacy_test/test_prod_op.py b/test/legacy_test/test_prod_op.py
@@ -19,59 +19,49 @@
 
 sys.path.append("../../legacy_test")
 from test_sum_op import TestReduceOPTensorAxisBase
+from utils import dygraph_guard, static_guard
 
 import paddle
+from paddle.framework import core
 
 
 class TestProdOp(unittest.TestCase):
     def setUp(self):
         self.input = np.random.random(size=(10, 10, 5)).astype(np.float32)
 
-    def run_imperative(self):
-        input = paddle.to_tensor(self.input)
-        dy_result = paddle.prod(input)
+    def run_imperative(self, place):
+        input = paddle.to_tensor(self.input, place=place)
+        out = paddle.prod(input)
         expected_result = np.prod(self.input)
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-        dy_result = paddle.prod(input, axis=1)
+        out = paddle.prod(input, axis=1)
         expected_result = np.prod(self.input, axis=1)
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-        dy_result = paddle.prod(input, axis=-1)
+        out = paddle.prod(input, axis=-1)
         expected_result = np.prod(self.input, axis=-1)
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-        dy_result = paddle.prod(input, axis=[0, 1])
+        out = paddle.prod(input, axis=[0, 1])
         expected_result = np.prod(self.input, axis=(0, 1))
         np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05, atol=1e-8
+            out.numpy(), expected_result, rtol=1e-05, atol=1e-8
         )
 
-        dy_result = paddle.prod(input, axis=1, keepdim=True)
+        out = paddle.prod(input, axis=1, keepdim=True)
         expected_result = np.prod(self.input, axis=1, keepdims=True)
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-        dy_result = paddle.prod(input, axis=1, dtype='int64')
+        out = paddle.prod(input, axis=1, dtype='int64')
         expected_result = np.prod(self.input, axis=1, dtype=np.int64)
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-        dy_result = paddle.prod(input, axis=1, keepdim=True, dtype='int64')
+        out = paddle.prod(input, axis=1, keepdim=True, dtype='int64')
         expected_result = np.prod(
             self.input, axis=1, keepdims=True, dtype=np.int64
         )
-        np.testing.assert_allclose(
-            dy_result.numpy(), expected_result, rtol=1e-05
-        )
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
     def run_static(self, use_gpu=False):
         with paddle.static.program_guard(paddle.static.Program()):
@@ -134,44 +124,135 @@ def run_static(self, use_gpu=False):
         )
 
     def test_cpu(self):
-        paddle.disable_static(place=paddle.CPUPlace())
-        self.run_imperative()
-        paddle.enable_static()
-
-        self.run_static()
+        with dygraph_guard():
+            self.run_imperative(place=paddle.CPUPlace())
+        with static_guard():
+            self.run_static()
 
     def test_gpu(self):
         if not paddle.base.core.is_compiled_with_cuda():
             return
+        with dygraph_guard():
+            self.run_imperative(place=paddle.CUDAPlace(0))
+        with static_guard():
+            self.run_static()
 
-        paddle.disable_static(place=paddle.CUDAPlace(0))
-        self.run_imperative()
-        paddle.enable_static()
 
-        self.run_static(use_gpu=True)
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestProdComplexOp(TestProdOp):
+    def setUp(self):
+        real = np.random.random(size=(10, 10, 5)).astype(np.float32)
+        imag = np.random.random(size=(10, 10, 5)).astype(np.float32)
+        self.input = real + 1j * imag
 
+    def run_imperative(self, place):
+        input = paddle.to_tensor(self.input, place=place)
+        out = paddle.prod(input)
+        expected_result = np.prod(self.input)
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-class TestProdOpError(unittest.TestCase):
+        out = paddle.prod(input, axis=1)
+        expected_result = np.prod(self.input, axis=1)
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
 
-    def test_error(self):
-        with paddle.static.program_guard(
-            paddle.static.Program(), paddle.static.Program()
-        ):
-            x = paddle.static.data(name='x', shape=[2, 2, 4], dtype='float32')
-            bool_x = paddle.static.data(
-                name='bool_x', shape=[2, 2, 4], dtype='bool'
+        out = paddle.prod(input, axis=[0, 1])
+        expected_result = np.prod(self.input, axis=(0, 1))
+        np.testing.assert_allclose(
+            out.numpy(), expected_result, rtol=1e-05, atol=1e-8
+        )
+
+        out = paddle.prod(input, axis=1, keepdim=True)
+        expected_result = np.prod(self.input, axis=1, keepdims=True)
+        np.testing.assert_allclose(out.numpy(), expected_result, rtol=1e-05)
+
+    def run_static(self, use_gpu=False):
+        with paddle.static.program_guard(paddle.static.Program()):
+            input = paddle.static.data(
+                name='input', shape=[10, 10, 5], dtype='complex64'
             )
-            # The argument x should be a Tensor
-            self.assertRaises(TypeError, paddle.prod, [1])
+            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)
 
-            # The data type of x should be float32, float64, int32, int64
-            self.assertRaises(TypeError, paddle.prod, bool_x)
+            place = paddle.CUDAPlace(0) if use_gpu else paddle.CPUPlace()
+            exe = paddle.static.Executor(place)
+            exe.run(paddle.static.default_startup_program())
+            static_complex_result = exe.run(
+                feed={"input": self.input},
+                fetch_list=[
+                    result0,
+                    result1,
+                    result2,
+                    result3,
+                    result4,
+                ],
+            )
+
+        expected_result = np.prod(self.input)
+        np.testing.assert_allclose(
+            static_complex_result[0], expected_result, rtol=1e-05
+        )
+        expected_result = np.prod(self.input, axis=1)
+        np.testing.assert_allclose(
+            static_complex_result[1], expected_result, rtol=1e-05
+        )
+        expected_result = np.prod(self.input, axis=-1)
+        np.testing.assert_allclose(
+            static_complex_result[2], expected_result, rtol=1e-05
+        )
+        expected_result = np.prod(self.input, axis=(0, 1))
+        np.testing.assert_allclose(
+            static_complex_result[3], expected_result, rtol=1e-05, atol=1e-8
+        )
+        expected_result = np.prod(self.input, axis=1, keepdims=True)
+        np.testing.assert_allclose(
+            static_complex_result[4], expected_result, rtol=1e-05
+        )
 
-            # The argument axis's type should be int ,list or tuple
-            self.assertRaises(TypeError, paddle.prod, x, 1.5)
+    def test_cpu(self):
+        with dygraph_guard():
+            self.run_imperative(place=paddle.CPUPlace())
+        with static_guard():
+            self.run_static()
 
-            # The argument dtype of prod_op should be float32, float64, int32 or int64.
-            self.assertRaises(TypeError, paddle.prod, x, 'bool')
+    def test_gpu(self):
+        if not paddle.base.core.is_compiled_with_cuda():
+            return
+        with dygraph_guard():
+            self.run_imperative(place=paddle.CUDAPlace(0))
+        with static_guard():
+            self.run_static()
+
+
+class TestProdOpError(unittest.TestCase):
+
+    def test_error(self):
+        with static_guard():
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                x = paddle.static.data(
+                    name='x', shape=[2, 2, 4], dtype='float32'
+                )
+                bool_x = paddle.static.data(
+                    name='bool_x', shape=[2, 2, 4], dtype='bool'
+                )
+                # The argument x should be a Tensor
+                self.assertRaises(TypeError, paddle.prod, [1])
+
+                # The data type of x should be float32, float64, int32, int64
+                self.assertRaises(TypeError, paddle.prod, bool_x)
+
+                # The argument axis's type should be int ,list or tuple
+                self.assertRaises(TypeError, paddle.prod, x, 1.5)
+
+                # The argument dtype of prod_op should be float32, float64, int32 or int64.
+                self.assertRaises(TypeError, paddle.prod, x, 'bool')
 
 
 class TestProdWithTensorAxis1(TestReduceOPTensorAxisBase):
