【Paddle Tensor No.10】新增 Tensor.__complex__ (#69257)

* add tensor complex

* update jit

* update convert_operators

* update test

* update test

* add test

Author: NKNaN

python/paddle/base/dygraph/math_op_patch.py

@@ -110,6 +110,16 @@ def _neg_(var: Tensor) -> Tensor:
     def _abs_(var: Tensor) -> Tensor:
         return var.abs()
 
+    def _complex_(var: Tensor) -> complex:
+        numel = np.prod(var.shape)
+        assert (
+            numel == 1
+        ), "only one element variable can be converted to complex."
+        assert var._is_initialized(), "variable's tensor is not initialized"
+        if not var.is_complex():
+            var = var.astype('complex64')
+        return complex(np.array(var))
+
     def _float_(var: Tensor) -> float:
         numel = np.prod(var.shape)
         assert (
@@ -203,6 +213,7 @@ def _mT_(var: Tensor) -> Tensor:
     eager_methods = [
         ('__neg__', _neg_),
         ('__abs__', _abs_),
+        ('__complex__', _complex_),
         ('__float__', _float_),
         ('__long__', _long_),
         ('__int__', _int_),

python/paddle/base/layers/math_op_patch.py

@@ -749,6 +749,13 @@ def _float_(self):
             "2. If you want to run it in full graph mode, you need use Variable directly, and do not use float(Variable)."
         )
 
+    def _complex_(self):
+        raise TypeError(
+            "complex(Variable) is not supported in static graph mode. If you are using @to_static, you can try this:\n"
+            "1. If you want to get the value of Variable, you can switch to non-fullgraph mode by setting @to_static(full_graph=True).\n"
+            "2. If you want to run it in full graph mode, you need use Variable directly, and do not use complex(Variable)."
+        )
+
     def values(var):
         block = current_block(var)
         out = create_new_tmp_var(block, var.dtype)
@@ -889,6 +896,7 @@ def to_dense(var):
         ('__ge__', _binary_creator_('__ge__', 'greater_equal', False, None)),
         ('__float__', _float_),
         ('__int__', _int_),
+        ('__complex__', _complex_),
         ('values', values),
         ('indices', indices),
         ('to_dense', to_dense),

python/paddle/jit/dy2static/convert_operators.py

@@ -759,18 +759,21 @@ def convert_var_dtype(var, dtype):
             'bool',
             'int',
             'float',
+            'complex',
         ], f"The casted target dtype is {dtype}, which is not supported in type casting."
         cast_map = {
             'bool': 'bool',
             'int': 'int32',
             'float': 'float32',
+            'complex': 'complex64',
         }
         return paddle.cast(var, dtype=cast_map[dtype])
     else:
         assert dtype in [
             'bool',
             'int',
             'float',
+            'complex',
         ], f"The casted target dtype is {dtype}, which is not supported in type casting."
         return eval(dtype)(var)
 

python/paddle/jit/dy2static/transformers/cast_transformer.py

@@ -27,7 +27,7 @@ class CastTransformer(BaseTransformer):
 
     def __init__(self, root):
         self.root = root
-        self._castable_type = {'bool', 'int', 'float'}
+        self._castable_type = {'bool', 'int', 'float', 'complex'}
 
     def transform(self):
         self.visit(self.root)

python/paddle/jit/sot/opcode_translator/executor/variables/basic.py

@@ -825,6 +825,12 @@ def __float__(self) -> float:
     def float(self):
         return ConstantVariable(float(self), self.graph, DummyTracker([self]))
 
+    def __complex__(self) -> complex:
+        return complex(self.get_py_value())
+
+    def complex(self):
+        return ConstantVariable(complex(self), self.graph, DummyTracker([self]))
+
     @property
     def out_var_name(self):
         return f"{self.graph.OUT_VAR_PREFIX}{self.var_name}"

python/paddle/jit/sot/utils/magic_methods.py

@@ -85,6 +85,7 @@
     float: "__float__",
     len: "__len__",
     int: "__int__",
+    complex: "__complex__",
 }
 # TODO(SigureMo): support any, all, sum
 

python/paddle/pir/math_op_patch.py

@@ -595,6 +595,30 @@ def _bool_(self):
             """
         raise TypeError(textwrap.dedent(error_msg))
 
+    def _complex_(self):
+        error_msg = """\
+            complex(Tensor) is not supported in static graph mode. Because it's value is not available during the static mode.
+            It's usually triggered by the logging implicitly, for example:
+                >>> logging.info("The value of x is: {complex(x)}")
+                                                              ^ `x` is Tensor, `complex(x)` triggers complex(Tensor)
+
+                There are two common workarounds available:
+                If you are logging Tensor values, then consider logging only at dynamic graphs, for example:
+
+                    Modify the following code
+                    >>> logging.info("The value of x is: {complex(x)}")
+                    to
+                    >>> if paddle.in_dynamic_mode():
+                    ...     logging.info("The value of x is: {complex(x)}")
+
+                If you need to convert the Tensor type, for example:
+                    Modify the following code
+                    >>> x = complex(x)
+                    to
+                    >>> x = x.astype("complex64")
+        """
+        raise TypeError(textwrap.dedent(error_msg))
+
     def clone(self):
         """
         Returns a new static Value, which is the clone of the original static
@@ -1143,6 +1167,7 @@ def register_hook(self, hook):
         ('__float__', _float_),
         ('__int__', _int_),
         ('__bool__', _bool_),
+        ('__complex__', _complex_),
     ]
 
     global _already_patch_value

python/paddle/tensor/tensor.prototype.pyi

@@ -187,6 +187,7 @@ class AbstractTensor:
     def __int__(self) -> int: ...
     def __long__(self) -> float: ...
     def __nonzero__(self) -> bool: ...
+    def __complex__(self) -> complex: ...
 
     # emulating container types
     def __getitem__(

test/dygraph_to_static/test_cast.py

@@ -18,6 +18,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     test_ast_only,
+    test_pir_only,
 )
 
 import paddle
@@ -58,6 +59,17 @@ def test_mix_cast(x):
     return x
 
 
+def test_complex_cast(x):
+    x = paddle.to_tensor(x)
+    x = complex(x)
+    return x
+
+
+def test_not_var_complex_cast(x):
+    x = complex(x)
+    return x
+
+
 class TestCastBase(Dy2StTestBase):
     def setUp(self):
         self.place = (
@@ -190,5 +202,64 @@ def test_cast_result(self):
         )
 
 
+@unittest.skipIf(
+    paddle.core.is_compiled_with_xpu(),
+    "xpu does not support complex cast temporarily",
+)
+class TestComplexCast(TestCastBase):
+    def prepare(self):
+        self.input_shape = (8, 16)
+        self.input_dtype = 'float32'
+        self.input = (
+            np.random.binomial(2, 0.5, size=np.prod(self.input_shape))
+            .reshape(self.input_shape)
+            .astype(self.input_dtype)
+        )
+        self.cast_dtype = 'complex64'
+
+    def set_func(self):
+        self.func = paddle.jit.to_static(full_graph=True)(test_complex_cast)
+
+    @test_pir_only
+    def test_cast_result(self):
+        self.set_func()
+        res = self.do_test().numpy()
+        self.assertTrue(
+            res.dtype == self.cast_dtype,
+            msg=f'The target dtype is {self.cast_dtype}, but the casted dtype is {res.dtype}.',
+        )
+        ref_val = self.input.astype(self.cast_dtype)
+        np.testing.assert_allclose(
+            res,
+            ref_val,
+            rtol=1e-05,
+            err_msg=f'The casted value is {res}.\nThe correct value is {ref_val}.',
+        )
+
+
+class TestNotVarComplexCast(TestCastBase):
+    def prepare(self):
+        self.input = 3.14
+        self.cast_dtype = 'complex'
+
+    def set_func(self):
+        self.func = paddle.jit.to_static(full_graph=True)(
+            test_not_var_complex_cast
+        )
+
+    @test_ast_only
+    def test_cast_result(self):
+        self.set_func()
+        res = self.do_test()
+        self.assertTrue(
+            type(res) == complex, msg='The casted dtype is not complex.'
+        )
+        ref_val = complex(self.input)
+        self.assertTrue(
+            res == ref_val,
+            msg=f'The casted value is {res}.\nThe correct value is {ref_val}.',
+        )
+
+
 if __name__ == '__main__':
     unittest.main()

test/legacy_test/test_math_op_patch.py

@@ -405,6 +405,8 @@ def test_builtin_type_conversion(self):
             int(a)
         with self.assertRaises(TypeError):
             float(a)
+        with self.assertRaises(TypeError):
+            complex(a)
 
 
 class TestDygraphMathOpPatches(unittest.TestCase):

test/legacy_test/test_math_op_patch_pir.py

@@ -734,6 +734,20 @@ def test_builtin_type_conversion(self):
                     float(x)
                 with self.assertRaises(TypeError):
                     bool(x)
+                with self.assertRaises(TypeError):
+                    complex(x)
+
+    def test_builtin_type_conversion_old_ir(self):
+        with paddle.pir_utils.DygraphOldIrGuard():
+            _, _, program_guard = new_program()
+            with program_guard:
+                x = paddle.static.data(name='x', shape=[], dtype="float32")
+                with self.assertRaises(TypeError):
+                    int(x)
+                with self.assertRaises(TypeError):
+                    float(x)
+                with self.assertRaises(TypeError):
+                    complex(x)
 
     def test_math_exists(self):
         with paddle.pir_utils.IrGuard():

test/legacy_test/test_math_op_patch_var_base.py

@@ -522,6 +522,14 @@ def test_float_int_long(self):
         self.assertTrue(int(a) == 999424)
         self.assertTrue(int(a) == 999424)
 
+    def test_complex(self):
+        with base.dygraph.guard():
+            a = paddle.to_tensor(np.array([100.1 + 99.9j]))
+            self.assertTrue(complex(a) == (100.1 + 99.9j))
+
+        a = paddle.to_tensor(1000000.0, dtype='bfloat16')
+        self.assertTrue(complex(a) == (999424 + 0j))
+
     def test_len(self):
         a_np = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
         with base.dygraph.guard():