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

fangfangssj · web-flow · commit e9d5c1fd3e12 · 2025-01-06T14:00:33.000+08:00
* sign

* add test

* fix

* fix

* fix

* fix
diff --git a/paddle/phi/kernels/cpu/sign_kernel.cc b/paddle/phi/kernels/cpu/sign_kernel.cc
@@ -31,4 +31,6 @@ PD_REGISTER_KERNEL(sign,
                    int32_t,
                    int64_t,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/funcs/eigen/sign.cc b/paddle/phi/kernels/funcs/eigen/sign.cc
@@ -11,6 +11,7 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
+#include "paddle/phi/common/complex.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
 
 namespace phi::funcs {
@@ -28,12 +29,39 @@ struct EigenSign<Eigen::DefaultDevice, T> {
   }
 };
 
+template <typename T>
+struct EigenSign<Eigen::DefaultDevice, phi::dtype::complex<T>> {
+  using InType = Eigen::TensorMap<Eigen::Tensor<const phi::dtype::complex<T>,
+                                                1,
+                                                Eigen::RowMajor,
+                                                Eigen::DenseIndex>>;
+  using OutType = Eigen::TensorMap<Eigen::Tensor<phi::dtype::complex<T>,
+                                                 1,
+                                                 Eigen::RowMajor,
+                                                 Eigen::DenseIndex>>;
+  static void Eval(const Eigen::DefaultDevice& dev,
+                   OutType out,
+                   const InType& in) {
+    out.device(dev) = in.unaryExpr(
+        [](const phi::dtype::complex<T>& z) -> phi::dtype::complex<T> {
+          T abs_val = abs(z);
+          if (abs_val == 0) {
+            return phi::dtype::complex<T>(0, 0);
+          } else {
+            return phi::dtype::complex<T>(z.real / abs_val, z.imag / abs_val);
+          }
+        });
+  }
+};
+
 template struct EigenSign<Eigen::DefaultDevice, uint8_t>;
 template struct EigenSign<Eigen::DefaultDevice, int8_t>;
 template struct EigenSign<Eigen::DefaultDevice, int16_t>;
 template struct EigenSign<Eigen::DefaultDevice, int32_t>;
 template struct EigenSign<Eigen::DefaultDevice, int64_t>;
 template struct EigenSign<Eigen::DefaultDevice, float>;
 template struct EigenSign<Eigen::DefaultDevice, double>;
+template struct EigenSign<Eigen::DefaultDevice, phi::dtype::complex<float>>;
+template struct EigenSign<Eigen::DefaultDevice, phi::dtype::complex<double>>;
 
 }  // namespace phi::funcs
diff --git a/paddle/phi/kernels/funcs/eigen/sign.cu b/paddle/phi/kernels/funcs/eigen/sign.cu
@@ -29,6 +29,30 @@ struct EigenSign<Eigen::GpuDevice, T> {
   }
 };
 
+template <typename T>
+struct EigenSign<Eigen::GpuDevice, phi::dtype::complex<T>> {
+  using InType = Eigen::TensorMap<Eigen::Tensor<const phi::dtype::complex<T>,
+                                                1,
+                                                Eigen::RowMajor,
+                                                Eigen::DenseIndex>>;
+  using OutType = Eigen::TensorMap<Eigen::Tensor<phi::dtype::complex<T>,
+                                                 1,
+                                                 Eigen::RowMajor,
+                                                 Eigen::DenseIndex>>;
+  static void Eval(const Eigen::GpuDevice& dev, OutType out, const InType& in) {
+    out.device(dev) = in.unaryExpr(
+        [] __host__ __device__(
+            const phi::dtype::complex<T>& z) -> phi::dtype::complex<T> {
+          T abs_val = abs(z);
+          if (abs_val == 0) {
+            return phi::dtype::complex<T>(0, 0);
+          } else {
+            return phi::dtype::complex<T>(z.real / abs_val, z.imag / abs_val);
+          }
+        });
+  }
+};
+
 template struct EigenSign<Eigen::GpuDevice, uint8_t>;
 template struct EigenSign<Eigen::GpuDevice, int8_t>;
 template struct EigenSign<Eigen::GpuDevice, int16_t>;
@@ -38,6 +62,8 @@ template struct EigenSign<Eigen::GpuDevice, float>;
 template struct EigenSign<Eigen::GpuDevice, double>;
 template struct EigenSign<Eigen::GpuDevice, dtype::float16>;
 template struct EigenSign<Eigen::GpuDevice, dtype::bfloat16>;
+template struct EigenSign<Eigen::GpuDevice, dtype::complex<float>>;
+template struct EigenSign<Eigen::GpuDevice, dtype::complex<double>>;
 
 }  // namespace funcs
 }  // namespace phi
diff --git a/paddle/phi/kernels/gpu/sign_kernel.cu.cc b/paddle/phi/kernels/gpu/sign_kernel.cu.cc
@@ -33,4 +33,6 @@ PD_REGISTER_KERNEL(sign,
                    float,
                    double,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4814,10 +4814,10 @@ def prod(
 
 def sign(x: Tensor, name: str | None = None) -> Tensor:
     """
-    Returns sign of every element in `x`: 1 for positive, -1 for negative and 0 for zero.
+    Returns sign of every element in `x`: For real numbers, 1 for positive, -1 for negative and 0 for zero. For complex numbers, the return value is a complex number with unit magnitude. If a complex number element is zero, the result is 0+0j.
 
     Args:
-        x (Tensor): The input tensor. The data type can be uint8, int8, int16, int32, int64, bfloat16, float16, float32 or float64.
+        x (Tensor): The input tensor. The data type can be uint8, int8, int16, int32, int64, bfloat16, float16, float32, float64, complex64 or complex128.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4850,6 +4850,8 @@ def sign(x: Tensor, name: str | None = None) -> Tensor:
                 'bfloat16',
                 'float32',
                 'float64',
+                'complex64',
+                'complex128',
             ],
             'sign',
         )
diff --git a/test/legacy_test/test_sign_op.py b/test/legacy_test/test_sign_op.py
@@ -25,6 +25,14 @@
 from paddle.base import core
 
 
+def complex_sign(x):
+    magnitude = np.abs(x)
+    result = np.zeros_like(x, dtype=x.dtype)
+    nonzero = magnitude != 0
+    result[nonzero] = x[nonzero] / magnitude[nonzero]
+    return result
+
+
 class TestSignOp(OpTest):
     def setUp(self):
         self.op_type = "sign"
@@ -51,6 +59,32 @@ def setUp(self):
         self.outputs = {'Out': np.sign(self.inputs['X'])}
 
 
+class TestSignComplex64Op(OpTest):
+    def setUp(self):
+        self.op_type = "sign"
+        self.python_api = paddle.sign
+        real_part = np.random.uniform(-10, 10, (10, 10))
+        imag_part = np.random.uniform(-10, 10, (10, 10))
+        self.inputs = {'X': (real_part + 1j * imag_part).astype("complex64")}
+        self.outputs = {'Out': complex_sign(self.inputs['X'])}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
+
+class TestSignComplex128Op(OpTest):
+    def setUp(self):
+        self.op_type = "sign"
+        self.python_api = paddle.sign
+        real_part = np.random.uniform(-10, 10, (10, 10))
+        imag_part = np.random.uniform(-10, 10, (10, 10))
+        self.inputs = {'X': (real_part + 1j * imag_part).astype("complex128")}
+        self.outputs = {'Out': complex_sign(self.inputs['X'])}
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_bfloat16_supported(core.CUDAPlace(0)),
@@ -167,6 +201,65 @@ def run(place):
             run(place)
 
 
+class TestSignComplexAPI(TestSignAPI):
+    def setUp(self):
+        self.place = []
+        self.place.append(base.CPUPlace())
+        if core.is_compiled_with_cuda():
+            self.place.append(base.CUDAPlace(0))
+
+    def test_dygraph(self):
+        with base.dygraph.guard():
+            np_x = np.array(
+                [-1.0 + 1.0j, 0.0 + 0.0j, 3.0 - 0.0j, 1.2 + 0.0j, 1.5 + 3.0j],
+                dtype='complex64',
+            )
+            x = paddle.to_tensor(np_x)
+            z = paddle.sign(x)
+            np_z = z.numpy()
+            z_expected = complex_sign(np_x)
+            np.testing.assert_allclose(np_z, z_expected, atol=1e-5, rtol=1e-5)
+
+    def test_static(self):
+        real_part = np.random.uniform(-10, 10, (10, 10))
+        imag_part = np.random.uniform(-10, 10, (10, 10))
+        np_input1 = (real_part + 1j * imag_part).astype("complex64")
+        np_input2 = (real_part + 1j * imag_part).astype("complex128")
+        np_out1 = complex_sign(np_input1)
+        np_out2 = complex_sign(np_input2)
+
+        def run(place):
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
+                # The input type of sign_op must be Variable or numpy.ndarray.
+                input1 = 12
+                self.assertRaises(TypeError, paddle.tensor.math.sign, input1)
+                # The result of sign_op must correct.
+                input1 = paddle.static.data(
+                    name='input1', shape=[12, 10], dtype="complex64"
+                )
+                input2 = paddle.static.data(
+                    name='input2', shape=[12, 10], dtype="complex128"
+                )
+                out1 = paddle.sign(input1)
+                out2 = paddle.sign(input2)
+                exe = paddle.static.Executor(place)
+                res1, res2 = exe.run(
+                    paddle.static.default_main_program(),
+                    feed={
+                        "input1": np_input1,
+                        "input2": np_input2,
+                    },
+                    fetch_list=[out1, out2],
+                )
+                np.testing.assert_allclose(res1, np_out1, atol=1e-5, rtol=1e-5)
+                np.testing.assert_allclose(res2, np_out2, atol=1e-5, rtol=1e-5)
+
+        for place in self.place:
+            run(place)
+
+
 class TestSignDoubleGradCheck(unittest.TestCase):
     def sign_wrapper(self, x):
         return paddle.sign(x[0])
