einsum support complex (#44212)

einsum support complex and add unittest.

paddle/phi/kernels/cpu/einsum_kernel.cc

@@ -18,5 +18,11 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/einsum_impl.h"
 
-PD_REGISTER_KERNEL(
-    einsum, CPU, ALL_LAYOUT, phi::EinsumKernelRaw, float, double) {}
+PD_REGISTER_KERNEL(einsum,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::EinsumKernelRaw,
+                   float,
+                   double,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/einsum_kernel.cu

@@ -25,4 +25,6 @@ PD_REGISTER_KERNEL(einsum,
                    float,
                    double,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}

python/paddle/fluid/tests/unittests/test_einsum_v2.py

@@ -541,5 +541,19 @@ class TestBF16(unittest.TestCase):
             self.assertEqual(C.item(), 8.0)
 
 
+class TestComplex(unittest.TestCase):
+    """
+    EinsumOp support Complex type
+    """
+
+    def test_shape(self):
+        a = paddle.rand([4, 4])
+        b = paddle.rand([4, 4])
+        c = paddle.einsum('xy,yz->xz', a, b)
+        a = paddle.cast(a, 'complex64')
+        b = paddle.cast(b, 'complex64')
+        c = paddle.einsum('xy,yz->xz', a, b)
+
+
 if __name__ == "__main__":
     unittest.main()