matrix_power support complex (#72308)

aquagull · web-flow · commit 555d7e3aabb5 · 2025-04-22T10:42:00.000+08:00
* update

* add test

* fix

* fix on xpu
diff --git a/paddle/phi/kernels/cpu/matrix_power_grad_kernel.cc b/paddle/phi/kernels/cpu/matrix_power_grad_kernel.cc
@@ -23,4 +23,6 @@ PD_REGISTER_KERNEL(matrix_power_grad,
                    ALL_LAYOUT,
                    phi::MatrixPowerGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/cpu/matrix_power_kernel.cc b/paddle/phi/kernels/cpu/matrix_power_kernel.cc
@@ -18,5 +18,11 @@ limitations under the License. */
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/matrix_power_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    matrix_power, CPU, ALL_LAYOUT, phi::MatrixPowerKernel, float, double) {}
+PD_REGISTER_KERNEL(matrix_power,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::MatrixPowerKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/matrix_power_grad_kernel.cu b/paddle/phi/kernels/gpu/matrix_power_grad_kernel.cu
@@ -23,4 +23,6 @@ PD_REGISTER_KERNEL(matrix_power_grad,
                    ALL_LAYOUT,
                    phi::MatrixPowerGradKernel,
                    float,
-                   double) {}
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/gpu/matrix_power_kernel.cu b/paddle/phi/kernels/gpu/matrix_power_kernel.cu
@@ -18,5 +18,11 @@ limitations under the License. */
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/matrix_power_kernel_impl.h"
 
-PD_REGISTER_KERNEL(
-    matrix_power, GPU, ALL_LAYOUT, phi::MatrixPowerKernel, float, double) {}
+PD_REGISTER_KERNEL(matrix_power,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::MatrixPowerKernel,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/test/legacy_test/test_matrix_power_op.py b/test/legacy_test/test_matrix_power_op.py
@@ -216,6 +216,122 @@ def config(self):
         self.n = 32
 
 
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpComplex64(TestMatrixPowerOp):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex64"
+        self.n = 2
+
+    def test_grad(self):
+        self.check_grad(["X"], "Out", max_relative_error=1e-2, check_pir=True)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpBatchedComplex64(TestMatrixPowerOpComplex64):
+    def config(self):
+        self.matrix_shape = [2, 8, 4, 4]
+        self.dtype = "complex64"
+        self.n = 2
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpLarge1Complex64(TestMatrixPowerOpComplex64):
+    def config(self):
+        self.matrix_shape = [32, 32]
+        self.dtype = "complex64"
+        self.n = 2
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpLarge2Complex64(TestMatrixPowerOpComplex64):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex64"
+        self.n = 32
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpComplex64Minus(TestMatrixPowerOpComplex64):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex64"
+        self.n = -1
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpComplex128(TestMatrixPowerOp):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex128"
+        self.n = 2
+
+    def test_grad(self):
+        self.check_grad(["X"], "Out", max_relative_error=1e-2, check_pir=True)
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpBatchedComplex128(TestMatrixPowerOpComplex128):
+    def config(self):
+        self.matrix_shape = [2, 8, 4, 4]
+        self.dtype = "complex128"
+        self.n = 2
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpLarge1Complex128(TestMatrixPowerOpComplex128):
+    def config(self):
+        self.matrix_shape = [32, 32]
+        self.dtype = "complex128"
+        self.n = 2
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpLarge2Complex128(TestMatrixPowerOpComplex128):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex128"
+        self.n = 32
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip complex due to lack of mean support",
+)
+class TestMatrixPowerOpComplex128Minus(TestMatrixPowerOpComplex128):
+    def config(self):
+        self.matrix_shape = [10, 10]
+        self.dtype = "complex128"
+        self.n = -1
+
+
 class TestMatrixPowerOpFP32(TestMatrixPowerOp):
     def config(self):
         self.matrix_shape = [10, 10]
