【Paddle Tensor 第二期 常用API复数类型支持 NO.6】 添加 full 函数复数类型支持 (#70277)

MrXnneHang · web-flow · commit 5859a053bbc0 · 2024-12-19T23:51:31.000+08:00
* 动态图正常，静态图编译成功，存挡.

* 静态图数值测试通过，需要确认dtype,存档.

* 验证静态图dtype.

* dtype修正完毕，准备丢到GPU测试。准备发射.

* 删除默认dtype的指定，增加对复数不输入dtype的测试覆盖

* 修复`bool-&gt;float32`的自动类型转换，添加覆盖测试.
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.cc b/paddle/fluid/pir/dialect/operator/ir/manual_api.cc
@@ -22,8 +22,10 @@
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/tensorrt_op.h"
+#include "paddle/phi/common/complex.h"
 #include "paddle/pir/include/core/builtin_op.h"
 #include "paddle/pir/include/core/parameter.h"
+
 namespace paddle::dialect {
 
 pir::Value builtin_combine(const std::vector<pir::Value>& x) {
@@ -54,6 +56,22 @@ std::vector<pir::Value> add_n_grad(const std::vector<pir::Value>& inputs,
   return inputs_grad;
 }
 
+pir::Value full(const std::vector<int64_t>& shape,
+                double real,
+                double imag,
+                phi::DataType dtype,
+                const phi::Place& place) {
+  CheckDataType(dtype, "dtype", "full");
+  if (dtype == phi::DataType::COMPLEX64) {
+    dtype = phi::DataType::FLOAT32;
+  } else {
+    dtype = phi::DataType::FLOAT64;
+  }
+  pir::Value real_tmp = full(shape, real, dtype, place);
+  pir::Value imag_tmp = full(shape, imag, dtype, place);
+  return paddle::dialect::complex(real_tmp, imag_tmp);
+}
+
 pir::Value zeros_like(const pir::Value& x,
                       const phi::DataType dtype,
                       const Place& place) {
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_api.h b/paddle/fluid/pir/dialect/operator/ir/manual_api.h
@@ -59,6 +59,12 @@ pir::Value ones(const std::vector<int64_t>& shape,
                 phi::DataType dtype = phi::DataType::FLOAT32,
                 const Place& place = phi::CPUPlace());
 
+pir::Value full(const std::vector<int64_t>& shape,
+                double real,
+                double imag,
+                phi::DataType dtype = phi::DataType::FLOAT32,
+                const Place& place = phi::CPUPlace());
+
 pir::Value ones_like(pir::Value x_,
                      phi::DataType dtype = phi::DataType::UNDEFINED,
                      const Place& place = {});
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -156,12 +156,23 @@ PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
         !PyObject_CheckIRVectorOfValue(shape_obj) &&
         !PyObject_CheckIRValue(value_obj)) {
       std::vector<int64_t> shape = CastPyArg2Longs(shape_obj, "full", 0);
-      double value = CastPyArg2Double(value_obj, "full", 1);
-      CallStackRecorder callstack_recoder("full");
-      callstack_recoder.Record();
-      auto static_api_out = paddle::dialect::full(shape, value, dtype, place);
-      callstack_recoder.AttachToOps();
-      return ToPyObject(static_api_out);
+      if (PyComplex_Check(value_obj)) {
+        phi::dtype::complex<float> complex_value =
+            CastPyArg2Complex(value_obj, "full", 1);
+        CallStackRecorder callstack_recoder("full");
+        callstack_recoder.Record();
+        auto static_api_out = paddle::dialect::full(
+            shape, complex_value.real, complex_value.imag, dtype, place);
+        callstack_recoder.AttachToOps();
+        return ToPyObject(static_api_out);
+      } else {
+        double value = CastPyArg2Double(value_obj, "full", 1);
+        CallStackRecorder callstack_recoder("full");
+        callstack_recoder.Record();
+        auto static_api_out = paddle::dialect::full(shape, value, dtype, place);
+        callstack_recoder.AttachToOps();
+        return ToPyObject(static_api_out);
+      }
     } else {
       pir::Value shape, value;
 
@@ -180,11 +191,21 @@ PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
       if (PyObject_CheckIRValue(value_obj)) {
         value = CastPyArg2Value(value_obj, "full", 1, false);
       } else {
-        double value_tmp = CastPyArg2Double(value_obj, "full", 1);
-        value = paddle::dialect::full(std::vector<int64_t>{1},
-                                      value_tmp,
-                                      phi::DataType::FLOAT32,
-                                      phi::CPUPlace());
+        if (PyComplex_Check(value_obj)) {
+          phi::dtype::complex<float> complex_value_tmp =
+              CastPyArg2Complex(value_obj, "full", 1);
+          value = paddle::dialect::full(std::vector<int64_t>{1},
+                                        complex_value_tmp.real,
+                                        complex_value_tmp.imag,
+                                        dtype,
+                                        place);
+        } else {
+          double value_tmp = CastPyArg2Double(value_obj, "full", 1);
+          value = paddle::dialect::full(std::vector<int64_t>{1},
+                                        value_tmp,
+                                        phi::DataType::FLOAT32,
+                                        phi::CPUPlace());
+        }
       }
 
       CallStackRecorder callstack_recoder("full_with_tensor");
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import builtins
 import math
 import re
 import warnings
@@ -1489,7 +1490,12 @@ def full(
     """
 
     if dtype is None:
-        dtype = paddle.get_default_dtype()
+        if isinstance(fill_value, (bool)):
+            dtype = "bool"
+        elif isinstance(fill_value, (builtins.complex)):
+            dtype = "complex128"
+        else:
+            dtype = paddle.get_default_dtype()
 
     return fill_constant(shape=shape, dtype=dtype, value=fill_value, name=name)
 
diff --git a/test/legacy_test/test_full_op.py b/test/legacy_test/test_full_op.py
@@ -67,8 +67,63 @@ def test_api(self):
             )
             out_8 = paddle.full(shape=10, dtype=np.float32, fill_value=val)
 
+            out_9 = paddle.full(
+                shape=10, dtype="complex64", fill_value=1.1 + 1.1j
+            )
+
+            out_10 = paddle.full(
+                shape=10, dtype="complex128", fill_value=1.1 + 1.1j
+            )
+
+            out_11 = paddle.full(
+                shape=10, dtype="complex64", fill_value=1.1 + np.inf * 1j
+            )
+
+            out_12 = paddle.full(
+                shape=10, dtype="complex128", fill_value=1.1 + np.inf * 1j
+            )
+
+            out_13 = paddle.full(
+                shape=10, dtype="complex64", fill_value=1.1 - np.inf * 1j
+            )
+
+            out_14 = paddle.full(
+                shape=10, dtype="complex128", fill_value=1.1 - np.inf * 1j
+            )
+
+            out_15 = paddle.full(
+                shape=10, dtype="complex64", fill_value=1.1 + np.nan * 1j
+            )
+
+            out_16 = paddle.full(
+                shape=10, dtype="complex128", fill_value=1.1 + np.nan * 1j
+            )
+
+            out_17 = paddle.full(shape=10, fill_value=1.1 + 1.1j)
+
+            out_18 = paddle.full(shape=10, fill_value=True)
+
             exe = base.Executor(place=base.CPUPlace())
-            res_1, res_2, res_3, res_4, res_5, res_6, res_7, res_8 = exe.run(
+            (
+                res_1,
+                res_2,
+                res_3,
+                res_4,
+                res_5,
+                res_6,
+                res_7,
+                res_8,
+                res_9,
+                res_10,
+                res_11,
+                res_12,
+                res_13,
+                res_14,
+                res_15,
+                res_16,
+                res_17,
+                res_18,
+            ) = exe.run(
                 paddle.static.default_main_program(),
                 feed={
                     "shape_tensor_int32": np.array([1, 2]).astype("int32"),
@@ -83,6 +138,16 @@ def test_api(self):
                     out_6,
                     out_7,
                     out_8,
+                    out_9,
+                    out_10,
+                    out_11,
+                    out_12,
+                    out_13,
+                    out_14,
+                    out_15,
+                    out_16,
+                    out_17,
+                    out_18,
                 ],
             )
 
@@ -110,6 +175,38 @@ def test_api(self):
         np.testing.assert_array_equal(
             res_8, np.full([10], 1.1, dtype="float32")
         )
+        np.testing.assert_allclose(
+            res_9, np.full([10], 1.1 + 1.1j, dtype="complex64")
+        )
+        np.testing.assert_allclose(
+            res_10, np.full([10], 1.1 + 1.1j, dtype="complex128")
+        )
+        np.testing.assert_allclose(
+            res_9, np.full([10], 1.1 + 1.1j, dtype="complex64")
+        )
+        np.testing.assert_allclose(
+            res_10, np.full([10], 1.1 + 1.1j, dtype="complex128")
+        )
+        np.testing.assert_allclose(
+            res_11, np.full([10], 1.1 + np.inf * 1j, dtype="complex64")
+        )
+        np.testing.assert_allclose(
+            res_12, np.full([10], 1.1 + np.inf * 1j, dtype="complex128")
+        )
+        np.testing.assert_allclose(
+            res_13, np.full([10], 1.1 - np.inf * 1j, dtype="complex64")
+        )
+        np.testing.assert_allclose(
+            res_14, np.full([10], 1.1 - np.inf * 1j, dtype="complex128")
+        )
+        np.testing.assert_allclose(
+            res_15, np.full([10], 1.1 + np.nan * 1j, dtype="complex64")
+        )
+        np.testing.assert_allclose(
+            res_16, np.full([10], 1.1 + np.nan * 1j, dtype="complex128")
+        )
+        np.testing.assert_allclose(res_17, np.full([10], 1.1 + 1.1j))
+        np.testing.assert_array_equal(res_18, np.full([10], True))
         paddle.disable_static()
 
     def test_api_eager(self):
@@ -166,6 +263,50 @@ def test_api_eager(self):
 
             out_11 = paddle.full(shape=10, dtype="float32", fill_value=1.1)
 
+            out_12 = paddle.full(
+                shape=[1, 2, 3], dtype="complex64", fill_value=1.1 + 1.1j
+            )
+
+            out_13 = paddle.full(
+                shape=[1, 2, 3], dtype="complex128", fill_value=1.1 + 1.1j
+            )
+
+            out_14 = paddle.full(
+                shape=[1, 2, 3], dtype="complex64", fill_value=1.1 + np.inf * 1j
+            )
+
+            out_15 = paddle.full(
+                shape=[1, 2, 3],
+                dtype="complex128",
+                fill_value=1.1 + np.inf * 1j,
+            )
+
+            out_16 = paddle.full(
+                shape=[1, 2, 3], dtype="complex64", fill_value=1.1 - np.inf * 1j
+            )
+
+            out_17 = paddle.full(
+                shape=[1, 2, 3],
+                dtype="complex128",
+                fill_value=1.1 - np.inf * 1j,
+            )
+
+            out_18 = paddle.full(
+                shape=[1, 2, 3], dtype="complex64", fill_value=1.1 + np.nan * 1j
+            )
+
+            out_19 = paddle.full(
+                shape=[1, 2, 3],
+                dtype="complex128",
+                fill_value=1.1 + np.nan * 1j,
+            )
+
+            # test without dtype input for complex
+            out_20 = paddle.full(shape=[1, 2, 3], fill_value=1.1 + 1.1j)
+
+            # test without dtype input for bool
+            out_21 = paddle.full(shape=[1, 2, 3], fill_value=True)
+
             np.testing.assert_array_equal(
                 out_1, np.full([1, 2], 1.1, dtype="float32")
             )
@@ -199,6 +340,35 @@ def test_api_eager(self):
             np.testing.assert_array_equal(
                 out_11, np.full([10], 1.1, dtype="float32")
             )
+            np.testing.assert_allclose(
+                out_12, np.full([1, 2, 3], 1.1 + 1.1j, dtype="complex64")
+            )
+            np.testing.assert_allclose(
+                out_13, np.full([1, 2, 3], 1.1 + 1.1j, dtype="complex128")
+            )
+            np.testing.assert_allclose(
+                out_14, np.full([1, 2, 3], 1.1 + np.inf * 1j, dtype="complex64")
+            )
+            np.testing.assert_allclose(
+                out_15,
+                np.full([1, 2, 3], 1.1 + np.inf * 1j, dtype="complex128"),
+            )
+            np.testing.assert_allclose(
+                out_16, np.full([1, 2, 3], 1.1 - np.inf * 1j, dtype="complex64")
+            )
+            np.testing.assert_allclose(
+                out_17,
+                np.full([1, 2, 3], 1.1 - np.inf * 1j, dtype="complex128"),
+            )
+            np.testing.assert_allclose(
+                out_18, np.full([1, 2, 3], 1.1 + np.nan * 1j, dtype="complex64")
+            )
+            np.testing.assert_allclose(
+                out_19,
+                np.full([1, 2, 3], 1.1 + np.nan * 1j, dtype="complex128"),
+            )
+            np.testing.assert_allclose(out_20, np.full([1, 2, 3], 1.1 + 1.1j))
+            np.testing.assert_array_equal(out_21, np.full([1, 2, 3], True))
 
 
 class TestFullOpError(unittest.TestCase):
