[Complex] Fix complex tensor print (PaddlePaddle#76380)

* fix complex tensor display

Author: HydrogenSulfate

python/paddle/tensor/to_string.py

@@ -136,6 +136,20 @@ def _format_item(np_var, max_width=0, signed=False):
             item_str = f'{np_var:.0f}.'
         else:
             item_str = f'{np_var:.{DEFAULT_PRINT_OPTIONS.precision}f}'
+    elif np_var.dtype == np.complex64 or np_var.dtype == np.complex128:
+        re = np.real(np_var)
+        im = np.imag(np_var)
+        prec = DEFAULT_PRINT_OPTIONS.precision
+        if DEFAULT_PRINT_OPTIONS.sci_mode:
+            if im >= 0:
+                item_str = f'({re:.{prec}e}+{im:.{prec}e}j)'
+            else:
+                item_str = f'({re:.{prec}e}{im:.{prec}e}j)'
+        else:
+            if im >= 0:
+                item_str = f'({re:.{prec}f}+{im:.{prec}f}j)'
+            else:
+                item_str = f'({re:.{prec}f}{im:.{prec}f}j)'
     else:
         item_str = f'{np_var}'
 
@@ -311,7 +325,9 @@ def _format_dense_tensor(tensor, indent):
     ):
         np_tensor = mask_xpu_bf16_tensor(np_tensor)
 
-    summary = tensor.numel() > DEFAULT_PRINT_OPTIONS.threshold
+    summary = (
+        np.prod(tensor.shape, dtype="int64") > DEFAULT_PRINT_OPTIONS.threshold
+    )
 
     max_width, signed = _get_max_width(_to_summary(np_tensor))
 

test/legacy_test/test_eager_tensor.py

@@ -24,6 +24,7 @@
 import paddle.nn.functional as F
 from paddle import base
 from paddle.base import core
+from paddle.tensor.to_string import DEFAULT_PRINT_OPTIONS
 from paddle.utils.dlpack import DLDeviceType
 
 
@@ -1337,6 +1338,72 @@ def test_tensor_str_fp8_e5m2(self):
 
         self.assertEqual(a_str, expected)
 
+    def test_tensor_str_complex64(self):
+        original_opt = copy.deepcopy(DEFAULT_PRINT_OPTIONS)
+        try:
+            paddle.disable_static(paddle.CPUPlace())
+            a = paddle.to_tensor(
+                [[1.5 + 1j, 1.0 - 2j], [0 - 3j, 0]], dtype="complex64"
+            ).cpu()
+            paddle.set_printoptions(precision=4)
+            a_str = str(a)
+
+            expected = """Tensor(shape=[2, 2], dtype=complex64, place=Place(cpu), stop_gradient=True,
+       [[(1.5000+1.0000j), (1.0000-2.0000j)],
+        [(0.0000-3.0000j), (0.0000+0.0000j)]])"""
+
+            self.assertEqual(a_str, expected)
+
+            paddle.set_printoptions(precision=4, sci_mode=True)
+            a_str = str(a)
+
+            expected = """Tensor(shape=[2, 2], dtype=complex64, place=Place(cpu), stop_gradient=True,
+       [[(1.5000e+00+1.0000e+00j), (1.0000e+00-2.0000e+00j)],
+        [(0.0000e+00-3.0000e+00j), (0.0000e+00+0.0000e+00j)]])"""
+
+            self.assertEqual(a_str, expected)
+        finally:
+            paddle.set_printoptions(
+                precision=original_opt.precision,
+                threshold=original_opt.threshold,
+                edgeitems=original_opt.edgeitems,
+                sci_mode=original_opt.sci_mode,
+                linewidth=original_opt.linewidth,
+            )
+
+    def test_tensor_str_complex128(self):
+        original_opt = copy.deepcopy(DEFAULT_PRINT_OPTIONS)
+        try:
+            paddle.disable_static(paddle.CPUPlace())
+            a = paddle.to_tensor(
+                [[1.5 + 1j, 1.0 - 2j], [0 - 3j, 0]], dtype="complex128"
+            ).cpu()
+            paddle.set_printoptions(precision=4)
+            a_str = str(a)
+
+            expected = """Tensor(shape=[2, 2], dtype=complex128, place=Place(cpu), stop_gradient=True,
+       [[(1.5000+1.0000j), (1.0000-2.0000j)],
+        [(0.0000-3.0000j), (0.0000+0.0000j)]])"""
+
+            self.assertEqual(a_str, expected)
+
+            paddle.set_printoptions(precision=4, sci_mode=True)
+            a_str = str(a)
+
+            expected = """Tensor(shape=[2, 2], dtype=complex128, place=Place(cpu), stop_gradient=True,
+       [[(1.5000e+00+1.0000e+00j), (1.0000e+00-2.0000e+00j)],
+        [(0.0000e+00-3.0000e+00j), (0.0000e+00+0.0000e+00j)]])"""
+
+            self.assertEqual(a_str, expected)
+        finally:
+            paddle.set_printoptions(
+                precision=original_opt.precision,
+                threshold=original_opt.threshold,
+                edgeitems=original_opt.edgeitems,
+                sci_mode=original_opt.sci_mode,
+                linewidth=original_opt.linewidth,
+            )
+
     def test_print_tensor_dtype(self):
         paddle.disable_static(paddle.CPUPlace())
         a = paddle.rand([1])