【complex op】No.19 add complex support for triangular_solve

paddle/phi/kernels/cpu/triangular_solve_grad_kernel.cc

@@ -20,4 +20,6 @@ PD_REGISTER_KERNEL(triangular_solve_grad,
                    ALL_LAYOUT,
                    phi::TriangularSolveGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/cpu/triangular_solve_kernel.cc

@@ -82,4 +82,6 @@ PD_REGISTER_KERNEL(triangular_solve,
                    ALL_LAYOUT,
                    phi::TriangularSolveKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/funcs/blas/blas_impl.h

@@ -877,7 +877,7 @@ struct CBlas<phi::dtype::complex<float>> {
                    const phi::dtype::complex<float> alpha,
                    const phi::dtype::complex<float> *A,
                    const int lda,
-                   phi::dtype::complex<double> *B,
+                   phi::dtype::complex<float> *B,
                    const int ldb) {
     cblas_ctrsm(layout, side, uplo, transA, diag, M, N, &alpha, A, lda, B, ldb);
   }

paddle/phi/kernels/funcs/matrix_reduce.cc

@@ -54,6 +54,8 @@ class MatrixReduceSumFunctor<T, CPUContext> {
 
 template class MatrixReduceSumFunctor<float, CPUContext>;
 template class MatrixReduceSumFunctor<double, CPUContext>;
+template class MatrixReduceSumFunctor<phi::dtype::complex<float>, CPUContext>;
+template class MatrixReduceSumFunctor<phi::dtype::complex<double>, CPUContext>;
 
 }  // namespace funcs
 }  // namespace phi

paddle/phi/kernels/funcs/matrix_reduce.cu

@@ -52,6 +52,8 @@ class MatrixReduceSumFunctor<T, GPUContext> {
 
 template class MatrixReduceSumFunctor<float, GPUContext>;
 template class MatrixReduceSumFunctor<double, GPUContext>;
+template class MatrixReduceSumFunctor<phi::dtype::complex<float>, GPUContext>;
+template class MatrixReduceSumFunctor<phi::dtype::complex<double>, GPUContext>;
 
 }  // namespace funcs
 }  // namespace phi

paddle/phi/kernels/gpu/triangular_solve_grad_kernel.cu

@@ -20,4 +20,6 @@ PD_REGISTER_KERNEL(triangular_solve_grad,
                    ALL_LAYOUT,
                    phi::TriangularSolveGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/triangular_solve_kernel.cu

@@ -128,4 +128,6 @@ PD_REGISTER_KERNEL(triangular_solve,
                    ALL_LAYOUT,
                    phi::TriangularSolveKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

python/paddle/tensor/linalg.py

@@ -3186,9 +3186,9 @@ def triangular_solve(
 
     Args:
         x (Tensor): The input triangular coefficient matrix. Its shape should be `[*, M, M]`, where `*` is zero or
-            more batch dimensions. Its data type should be float32 or float64.
+            more batch dimensions. Its data type should be float32, float64, complex64, complex128.
         y (Tensor): Multiple right-hand sides of system of equations. Its shape should be `[*, M, K]`, where `*` is
-            zero or more batch dimensions. Its data type should be float32 or float64.
+            zero or more batch dimensions. Its data type should be float32, float64, complex64, complex128.
         upper (bool, optional): Whether to solve the upper-triangular system of equations (default) or the lower-triangular
             system of equations. Default: True.
         transpose (bool, optional): whether `x` should be transposed before calculation. Default: False.
@@ -3227,10 +3227,16 @@ def triangular_solve(
         inputs = {"X": [x], "Y": [y]}
         helper = LayerHelper("triangular_solve", **locals())
         check_variable_and_dtype(
-            x, 'x', ['float32', 'float64'], 'triangular_solve'
+            x,
+            'x',
+            ['float32', 'float64', 'complex64', 'complex128'],
+            'triangular_solve',
         )
         check_variable_and_dtype(
-            y, 'y', ['float32', 'float64'], 'triangular_solve'
+            y,
+            'y',
+            ['float32', 'float64', 'complex64', 'complex128'],
+            'triangular_solve',
         )
         out = helper.create_variable_for_type_inference(dtype=x.dtype)
 

test/legacy_test/test_triangular_solve_op.py

@@ -51,10 +51,23 @@ def setUp(self):
         self.python_api = paddle.tensor.linalg.triangular_solve
         self.config()
 
-        self.inputs = {
-            'X': np.random.random(self.x_shape).astype(self.dtype),
-            'Y': np.random.random(self.y_shape).astype(self.dtype),
-        }
+        if self.dtype is np.complex64 or self.dtype is np.complex128:
+            self.inputs = {
+                'X': (
+                    np.random.random(self.x_shape)
+                    + 1j * np.random.random(self.x_shape)
+                ).astype(self.dtype),
+                'Y': (
+                    np.random.random(self.y_shape)
+                    + 1j * np.random.random(self.y_shape)
+                ).astype(self.dtype),
+            }
+        else:
+            self.inputs = {
+                'X': np.random.random(self.x_shape).astype(self.dtype),
+                'Y': np.random.random(self.y_shape).astype(self.dtype),
+            }
+
         self.attrs = {
             'upper': self.upper,
             'transpose': self.transpose,
@@ -248,6 +261,46 @@ def set_output(self):
         self.output = np.matmul(np.linalg.inv(x), y)
 
 
+# 3D(broadcast) + 3D complex64
+class TestTriangularSolveOpCp64(TestTriangularSolveOp):
+    """
+    case complex64
+    """
+
+    def config(self):
+        self.x_shape = [1, 10, 10]
+        self.y_shape = [6, 10, 12]
+        self.upper = False
+        self.transpose = False
+        self.unitriangular = False
+        self.dtype = "complex64"
+
+    def set_output(self):
+        x = np.tril(self.inputs['X'])
+        y = self.inputs['Y']
+        self.output = np.linalg.solve(x, y)
+
+
+# 3D(broadcast) + 3D complex128
+class TestTriangularSolveCp128(TestTriangularSolveOp):
+    """
+    case complex128
+    """
+
+    def config(self):
+        self.x_shape = [1, 10, 10]
+        self.y_shape = [6, 10, 12]
+        self.upper = False
+        self.transpose = False
+        self.unitriangular = False
+        self.dtype = "complex128"
+
+    def set_output(self):
+        x = np.tril(self.inputs['X'])
+        y = self.inputs['Y']
+        self.output = np.linalg.solve(x, y)
+
+
 class TestTriangularSolveAPI(unittest.TestCase):
     def setUp(self):
         np.random.seed(2021)