From 2aedd9db9684142e6338f8ae2d111da216f3f506 Mon Sep 17 00:00:00 2001
From: zhwesky2010 <1183042833@qq.com>
Date: Mon, 8 May 2023 15:12:59 +0800
Subject: [PATCH] [BUG] fix paddle.to_tensor/Tensor.item/Tensor.numpy BF16 bug
 (#53567)

---
 paddle/fluid/pybind/eager_method.cc           |  4 +-
 python/paddle/fluid/data_feeder.py            | 23 +++++++----
 .../fluid/dygraph/tensor_patch_methods.py     | 10 ++++-
 .../fluid/tests/unittests/test_var_base.py    | 26 +++++++++++-
 python/paddle/framework/framework.py          |  4 +-
 python/paddle/hapi/progressbar.py             |  2 -
 python/paddle/tensor/creation.py              | 41 +++++++++++--------
 7 files changed, 73 insertions(+), 37 deletions(-)

diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
index be2771293d40d..2caf978db6f23 100644
--- a/paddle/fluid/pybind/eager_method.cc
+++ b/paddle/fluid/pybind/eager_method.cc
@@ -1079,13 +1079,11 @@ static PyObject* tensor__getitem_from_offset(TensorObject* self,
     T b = paddle::pybind::TensorGetElement<T>(tensor, offset);               \
     Py_intptr_t py_dims[paddle::framework::DDim::kMaxRank];                  \
     Py_intptr_t py_strides[paddle::framework::DDim::kMaxRank];               \
-    py_dims[0] = 1;                                                          \
-    py_strides[0] = 1;                                                       \
     auto& api = pybind11::detail::npy_api::get();                            \
     PyObject* array = api.PyArray_NewFromDescr_(                             \
         api.PyArray_Type_,                                                   \
         api.PyArray_DescrFromType_(numpy_dtype),                             \
-        1,                                                                   \
+        0,                                                                   \
         py_dims,                                                             \
         py_strides,                                                          \
         nullptr,                                                             \
diff --git a/python/paddle/fluid/data_feeder.py b/python/paddle/fluid/data_feeder.py
index 77abddce03d82..908fedf7c4b9c 100644
--- a/python/paddle/fluid/data_feeder.py
+++ b/python/paddle/fluid/data_feeder.py
@@ -46,10 +46,6 @@
 }
 
 
-def copy_bits_from_float_to_uint16(f):
-    return struct.unpack('<I', struct.pack('<f', f))[0] >> 16
-
-
 def convert_float_to_uint16(data, data_format="NCHW"):
     if data.size == 0:
         return data.view(np.uint16)
@@ -57,16 +53,25 @@ def convert_float_to_uint16(data, data_format="NCHW"):
     if data_format == "NHWC":
         data = np.transpose(data, [0, 3, 1, 2])
 
-    new_data = []
-    for x in np.nditer(data):
-        new_data.append(np.uint16(copy_bits_from_float_to_uint16(x)))
-    new_data = np.reshape(new_data, data.shape).view(np.uint16)
+    new_data = np.vectorize(
+        lambda x: struct.unpack('<I', struct.pack('<f', x))[0] >> 16,
+        otypes=[np.uint16],
+    )(data.flat)
+    new_data = np.reshape(new_data, data.shape)
 
     if data_format == "NHWC":
-        new_data = np.transpose(new_output, [0, 2, 3, 1])
+        new_data = np.transpose(new_data, [0, 2, 3, 1])
     return new_data
 
 
+def convert_uint16_to_float(data):
+    new_data = np.vectorize(
+        lambda x: struct.unpack('<f', struct.pack('<I', x << 16))[0],
+        otypes=[np.float32],
+    )(data.flat)
+    return np.reshape(new_data, data.shape)
+
+
 def convert_dtype(dtype):
     if isinstance(dtype, core.VarDesc.VarType):
         if dtype in _PADDLE_DTYPE_2_NUMPY_DTYPE:
diff --git a/python/paddle/fluid/dygraph/tensor_patch_methods.py b/python/paddle/fluid/dygraph/tensor_patch_methods.py
index 90d896f7c7dc1..80d2b5e973cfe 100644
--- a/python/paddle/fluid/dygraph/tensor_patch_methods.py
+++ b/python/paddle/fluid/dygraph/tensor_patch_methods.py
@@ -33,7 +33,10 @@
 )
 from .base import switch_to_static_graph
 from .math_op_patch import monkey_patch_math_tensor
-from paddle.fluid.data_feeder import convert_dtype, _PADDLE_DTYPE_2_NUMPY_DTYPE
+from paddle.fluid.data_feeder import (
+    convert_uint16_to_float,
+    _PADDLE_DTYPE_2_NUMPY_DTYPE,
+)
 import paddle.utils.deprecated as deprecated
 import paddle.profiler as profiler
 from paddle.profiler.utils import in_profiler_mode
@@ -614,7 +617,10 @@ def item(self, *args):
                 print(x.item(0, 2))         #3.3
 
         """
-        return self._getitem_from_offset(*args).item()
+        scalar = self._getitem_from_offset(*args)
+        if scalar.dtype == np.uint16:
+            return convert_uint16_to_float(scalar).item()
+        return scalar.item()
 
     @property
     def inplace_version(self):
diff --git a/python/paddle/fluid/tests/unittests/test_var_base.py b/python/paddle/fluid/tests/unittests/test_var_base.py
index 9d489c77374c2..5fc9914fa1e88 100644
--- a/python/paddle/fluid/tests/unittests/test_var_base.py
+++ b/python/paddle/fluid/tests/unittests/test_var_base.py
@@ -246,10 +246,12 @@ def check_with_place(place):
                 np.testing.assert_array_equal(x.numpy(), numpy_array)
                 self.assertEqual(x.type, core.VarDesc.VarType.LOD_TENSOR)
 
-                # test dtype bfloat16
+                # test dtype=bfloat16
                 x = paddle.to_tensor(-1e6, dtype=paddle.bfloat16)
                 self.assertEqual(x.dtype, core.VarDesc.VarType.BF16)
                 self.assertTrue(x == -999424.0)
+                self.assertTrue(x.item() == -999424.0)
+                self.assertTrue(isinstance(x.item(), float))
 
                 x = paddle.to_tensor([-1e6, -1e6, -1e6], dtype='bfloat16')
                 self.assertEqual(x.dtype, core.VarDesc.VarType.BF16)
@@ -266,6 +268,28 @@ def check_with_place(place):
                 y.backward()
                 self.assertTrue(x.grad == -999424.0 * 2)
 
+                # test default_type=bfloat16
+                paddle.set_default_dtype('bfloat16')
+                x = paddle.to_tensor(-1e6)
+                self.assertEqual(x.dtype, core.VarDesc.VarType.BF16)
+                self.assertTrue(x == -999424.0)
+                self.assertTrue(x.item() == -999424.0)
+                self.assertTrue(isinstance(x.item(), float))
+
+                x = paddle.to_tensor([-1e6, -1e6, -1e6])
+                self.assertEqual(x.dtype, core.VarDesc.VarType.BF16)
+                self.assertTrue(x[0] == -999424.0)
+                self.assertTrue(x[1] == -999424.0)
+                self.assertTrue(x[2] == -999424.0)
+
+                x = paddle.to_tensor(-1e6, stop_gradient=False)
+                self.assertEqual(x.dtype, core.VarDesc.VarType.BF16)
+                self.assertTrue(x == -999424.0)
+                y = x * x
+                y.backward()
+                self.assertTrue(x.grad == -999424.0 * 2)
+                paddle.set_default_dtype('float32')
+
                 with self.assertRaises(ValueError):
                     paddle.randn([3, 2, 2]).item()
                 with self.assertRaises(ValueError):
diff --git a/python/paddle/framework/framework.py b/python/paddle/framework/framework.py
index 2e9df1a8ee353..563c4b4cd3b20 100644
--- a/python/paddle/framework/framework.py
+++ b/python/paddle/framework/framework.py
@@ -40,7 +40,7 @@ def set_default_dtype(d):
 
     """
     if isinstance(d, type):
-        # This branch is for NumPy scalar types
+        # This branch is for np.dtype
         if d in [np.float16, np.float32, np.float64]:
             d = d.__name__
         else:
@@ -49,7 +49,7 @@ def set_default_dtype(d):
                 ", but received %s" % d.__name__
             )
     else:
-        # This branch is for np.dtype and str
+        # This branch is for str
         if d in ['float16', 'float32', 'float64', 'bfloat16']:
             # NOTE(SigureMo): Since the np.dtype object is not an instance of
             # type, so it will not be handled by the previous branch. We need
diff --git a/python/paddle/hapi/progressbar.py b/python/paddle/hapi/progressbar.py
index 08d6cf2b12181..c5b25a58c12e1 100644
--- a/python/paddle/hapi/progressbar.py
+++ b/python/paddle/hapi/progressbar.py
@@ -83,8 +83,6 @@ def convert_uint16_to_float(in_list):
             if k == "loss":
                 if isinstance(val, list):
                     scalar_val = val[0]
-                elif isinstance(val, np.ndarray):
-                    scalar_val = val.item()
                 else:
                     scalar_val = val
                 if isinstance(scalar_val, np.uint16):
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
index db1006b0b1af0..0dca909bd5833 100644
--- a/python/paddle/tensor/creation.py
+++ b/python/paddle/tensor/creation.py
@@ -542,18 +542,28 @@ def logspace(start, stop, num, base=10.0, dtype=None, name=None):
 
 
 def _to_tensor_non_static(data, dtype=None, place=None, stop_gradient=True):
+    def _handle_tensor_dtype(tensor, dtype):
+        if dtype:
+            if convert_dtype(dtype) != convert_dtype(tensor.dtype):
+                return tensor.astype(convert_dtype(dtype))
+        return tensor
+
+    def _handle_np_dtype(ndarray, dtype):
+        if dtype:
+            if convert_dtype(dtype) != convert_dtype(ndarray.dtype):
+                # should not ndarray.astype('uint16') directly, data bits is wrong
+                if convert_dtype(dtype) in ['uint16']:
+                    return convert_float_to_uint16(ndarray.astype('float32'))
+                else:
+                    return ndarray.astype(convert_dtype(dtype))
+
+        return ndarray
 
     if isinstance(data, np.number):  # Special case for numpy scalars
         data = np.array(data)
 
     if not isinstance(data, np.ndarray):
 
-        def _handle_dtype(data, dtype):
-            if dtype:
-                if convert_dtype(dtype) != convert_dtype(data.dtype):
-                    return data.astype(convert_dtype(dtype))
-            return data
-
         if np.isscalar(data) and not isinstance(data, str):
             data = np.array(data)
         elif isinstance(data, (list, tuple)):
@@ -565,12 +575,12 @@ def _handle_dtype(data, dtype):
                 )
         elif isinstance(data, paddle.Tensor) and not in_dygraph_mode():
             data = data._copy_to(place, False)
-            data = _handle_dtype(data, dtype)
+            data = _handle_tensor_dtype(data, dtype)
             data.stop_gradient = stop_gradient
             return data
         elif isinstance(data, core.eager.Tensor) and in_dygraph_mode():
             data = data._copy_to(place, False)
-            data = _handle_dtype(data, dtype)
+            data = _handle_tensor_dtype(data, dtype)
             data.stop_gradient = stop_gradient
             return data
         elif isinstance(data, (core.LoDTensor, core.Tensor)):
@@ -583,7 +593,7 @@ def _handle_dtype(data, dtype):
                 data = paddle.Tensor(data)
             if not data.place._equals(place):
                 data = data._copy_to(place, False)
-            data = _handle_dtype(data, dtype)
+            data = _handle_tensor_dtype(data, dtype)
             data.stop_gradient = stop_gradient
             return data
         else:
@@ -607,18 +617,13 @@ def _handle_dtype(data, dtype):
                         if default_type in ['float16', 'float32']
                         else 'complex128'
                     )
-                data = data.astype(default_type)
+                data = _handle_np_dtype(data, default_type)
             # Windows default type is 'int32', while Linux/Mac is 'int64'. Unify they.
             if data.dtype in ['int32']:
-                default_type = "int64"
-                data = data.astype(default_type)
+                data = data.astype("int64")
 
-    if dtype and convert_dtype(dtype) != data.dtype:
-        if convert_dtype(dtype) in ['uint16']:
-            # should not ndarray.astype('uint16') directly, data bits is wrong
-            data = convert_float_to_uint16(data.astype('float32'))
-        else:
-            data = data.astype(convert_dtype(dtype))
+    if dtype:
+        data = _handle_np_dtype(data, dtype)
 
     if _in_eager_without_dygraph_check() and isinstance(data, np.ndarray):
         return core.eager.Tensor(