【complex op】 add complex support for tril,triu and tril_triu (#58532)

zyt1024 · web-flow · commit 740e4f8d37b4 · 2023-11-09T10:52:12.000+08:00
* [complex] add complex support for tril,triu and tril_triu

* complex tril_triu for xpu

* [complex] add complex support for tril,triu and tril_triu

* [complex] add complex support for tril,triu and tril_triu only cpu and gpu
diff --git a/paddle/phi/kernels/cpu/tril_triu_grad_kernel.cc b/paddle/phi/kernels/cpu/tril_triu_grad_kernel.cc
@@ -25,7 +25,9 @@ PD_REGISTER_KERNEL(tril_triu_grad,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(triu_grad,
                    CPU,
@@ -36,7 +38,9 @@ PD_REGISTER_KERNEL(triu_grad,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(tril_grad,
                    CPU,
@@ -47,4 +51,6 @@ PD_REGISTER_KERNEL(tril_grad,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/tril_triu_kernel.cc b/paddle/phi/kernels/cpu/tril_triu_kernel.cc
@@ -25,7 +25,9 @@ PD_REGISTER_KERNEL(tril_triu,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(triu,
                    CPU,
@@ -36,7 +38,9 @@ PD_REGISTER_KERNEL(triu,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(tril,
                    CPU,
@@ -47,4 +51,6 @@ PD_REGISTER_KERNEL(tril,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/tril_triu_grad_kernel.cu b/paddle/phi/kernels/gpu/tril_triu_grad_kernel.cu
@@ -26,7 +26,9 @@ PD_REGISTER_KERNEL(tril_grad,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(triu_grad,
                    GPU,
@@ -38,7 +40,9 @@ PD_REGISTER_KERNEL(triu_grad,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(tril_triu_grad,
                    GPU,
@@ -50,4 +54,6 @@ PD_REGISTER_KERNEL(tril_triu_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/gpu/tril_triu_kernel.cu b/paddle/phi/kernels/gpu/tril_triu_kernel.cu
@@ -26,7 +26,9 @@ PD_REGISTER_KERNEL(tril_triu,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(triu,
                    GPU,
@@ -38,7 +40,9 @@ PD_REGISTER_KERNEL(triu,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(tril,
                    GPU,
@@ -50,4 +54,6 @@ PD_REGISTER_KERNEL(tril,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -1435,7 +1435,17 @@ def _tril_triu_op(helper):
     check_variable_and_dtype(
         x,
         'x',
-        ['float16', 'uint16', 'float32', 'float64', 'int32', 'int64', 'bool'],
+        [
+            'float16',
+            'uint16',
+            'float32',
+            'float64',
+            'int32',
+            'int64',
+            'bool',
+            'complex64',
+            'complex128',
+        ],
         op_type,
     )
     if len(x.shape) < 2:
@@ -1474,7 +1484,7 @@ def tril(x, diagonal=0, name=None):
 
     Args:
         x (Tensor): The input x which is a Tensor.
-            Support data types: ``bool``, ``float64``, ``float32``, ``int32``, ``int64``.
+            Support data types: ``bool``, ``float64``, ``float32``, ``int32``, ``int64``, ``complex64``, ``complex128``.
         diagonal (int, optional): The diagonal to consider, default value is 0.
             If :attr:`diagonal` = 0, all elements on and below the main diagonal are
             retained. A positive value includes just as many diagonals above the main
@@ -1549,7 +1559,7 @@ def triu(x, diagonal=0, name=None):
 
     Args:
         x (Tensor): The input x which is a Tensor.
-            Support data types: ``float64``, ``float32``, ``int32``, ``int64``.
+            Support data types: ``float64``, ``float32``, ``int32``, ``int64``, ``complex64``, ``complex128``.
         diagonal (int, optional): The diagonal to consider, default value is 0.
             If :attr:`diagonal` = 0, all elements on and above the main diagonal are
             retained. A positive value excludes just as many diagonals above the main
diff --git a/test/legacy_test/test_tril_triu_op.py b/test/legacy_test/test_tril_triu_op.py
@@ -58,13 +58,28 @@ def initTestCase(self):
         self.real_op_type = np.random.choice(['triu', 'tril'])
         self.diagonal = None
         self.X = np.arange(1, 101, dtype=self.dtype).reshape([10, -1])
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            self.X = (
+                np.random.uniform(-1, 1, [10, 10])
+                + 1j * np.random.uniform(-1, 1, [10, 10])
+            ).astype(self.dtype)
 
 
 class TrilTriuOpDefaultTestFP16(TrilTriuOpDefaultTest):
     def init_dtype(self):
         self.dtype = np.float16
 
 
+class TrilTriuOpDefaultTestComplex_64(TrilTriuOpDefaultTest):
+    def init_dtype(self):
+        self.dtype = np.complex64
+
+
+class TrilTriuOpDefaultTestComplex_128(TrilTriuOpDefaultTest):
+    def init_dtype(self):
+        self.dtype = np.complex128
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_bfloat16_supported(core.CUDAPlace(0)),
@@ -146,6 +161,24 @@ def initTestCase(self):
             self.diagonal = diagonal
             self.X = np.random.random(Xshape).astype("float32")
 
+    class SuccessCaseComplex64(TrilTriuOpDefaultTestComplex_64):
+        def initTestCase(self):
+            self.init_dtype()
+            self.real_op_type = op_type
+            self.diagonal = diagonal
+            self.X = (
+                np.random.random(Xshape) + 1j * np.random.random(Xshape)
+            ).astype("complex64")
+
+    class SuccessCaseComplex128(TrilTriuOpDefaultTestComplex_128):
+        def initTestCase(self):
+            self.init_dtype()
+            self.real_op_type = op_type
+            self.diagonal = diagonal
+            self.X = (
+                np.random.random(Xshape) + 1j * np.random.random(Xshape)
+            ).astype("complex128")
+
     if dtype == "float64":
         CLASS = locals()[
             'SuccessCase' if expected == "success" else 'FailureCase'
@@ -158,6 +191,14 @@ def initTestCase(self):
         CLASS = locals()[
             'SuccessCaseBF16' if expected == "success" else 'FailureCase'
         ]
+    elif dtype == "complex64":
+        CLASS = locals()[
+            'SuccessCaseComplex64' if expected == "success" else 'FailureCase'
+        ]
+    elif dtype == "complex128":
+        CLASS = locals()[
+            'SuccessCaseComplex128' if expected == "success" else 'FailureCase'
+        ]
     else:
         raise ValueError(f"Not supported dtype {dtype}")
     CLASS.__name__ = cls_name
@@ -185,7 +226,7 @@ def initTestCase(self):
         (2020,): [None],
     },
 }
-for dtype in ["float64", "float16", "bfloat16"]:
+for dtype in ["float64", "float16", "bfloat16", "complex64", "complex128"]:
     for _op_type in ['tril', 'triu']:
         for _expected, _params in cases.items():
             for _Xshape, _diaglist in _params.items():
@@ -204,7 +245,7 @@ class TestTrilTriuOpAPI(unittest.TestCase):
     def test_api(self):
         paddle.enable_static()
 
-        dtypes = ['float16', 'float32']
+        dtypes = ['float16', 'float32', 'complex64', 'complex128']
         for dtype in dtypes:
             prog = paddle.static.Program()
             startup_prog = paddle.static.Program()
@@ -213,6 +254,11 @@ def test_api(self):
                 x = paddle.static.data(
                     shape=[1, 9, -1, 4], dtype=dtype, name='x'
                 )
+                if dtype == 'complex64' or dtype == 'complex128':
+                    data = (
+                        np.random.uniform(-1, 1, [1, 9, 9, 4])
+                        + 1j * np.random.uniform(-1, 1, [1, 9, 9, 4])
+                    ).astype(dtype)
                 tril_out, triu_out = tensor.tril(x), tensor.triu(x)
 
                 place = (
@@ -232,10 +278,15 @@ def test_api(self):
     def test_api_with_dygraph(self):
         paddle.disable_static()
 
-        dtypes = ['float16', 'float32']
+        dtypes = ['float16', 'float32', 'complex64', 'complex128']
         for dtype in dtypes:
             with base.dygraph.guard():
                 data = np.random.random([1, 9, 9, 4]).astype(dtype)
+                if dtype == 'complex64' or dtype == 'complex128':
+                    data = (
+                        np.random.uniform(-1, 1, [1, 9, 9, 4])
+                        + 1j * np.random.uniform(-1, 1, [1, 9, 9, 4])
+                    ).astype(dtype)
                 x = base.dygraph.to_variable(data)
                 tril_out, triu_out = (
                     tensor.tril(x).numpy(),
@@ -248,7 +299,7 @@ def test_api_with_dygraph(self):
     def test_base_api(self):
         paddle.enable_static()
 
-        dtypes = ['float16', 'float32']
+        dtypes = ['float16', 'float32', 'complex64', 'complex128']
         for dtype in dtypes:
             prog = paddle.static.Program()
             startup_prog = paddle.static.Program()
@@ -257,6 +308,11 @@ def test_base_api(self):
                 x = paddle.static.data(
                     shape=[1, 9, -1, 4], dtype=dtype, name='x'
                 )
+                if dtype == 'complex64' or dtype == 'complex128':
+                    data = (
+                        np.random.uniform(-1, 1, [1, 9, 9, 4])
+                        + 1j * np.random.uniform(-1, 1, [1, 9, 9, 4])
+                    ).astype(dtype)
                 triu_out = paddle.triu(x)
 
                 place = (
