Add fp16 support in squeeze & assign api (#22912)

test=release/1.7

python/paddle/fluid/layers/nn.py

@@ -5712,7 +5712,7 @@ def squeeze(input, axes, name=None):
             Out.shape = [1,3,5]
 
     Args:
-        input (Variable): The input Tensor. Support data type: float32, float64, int8, int32, int64.
+        input (Variable): The input Tensor. Support data type: float16, float32, float64, int8, int32, int64.
                           axes (list): One integer or List of integers, indicating the dimensions to be squeezed.
                           Axes range is :math:`[-rank(input), rank(input))`.
                           If axes is negative, :math:`axes=axes+rank(input)`.
@@ -5732,9 +5732,9 @@ def squeeze(input, axes, name=None):
 
     """
     helper = LayerHelper("squeeze", **locals())
-    check_type_and_dtype(input, 'input', Variable,
-                         ['float32', 'float64', 'int8', 'int32', 'int64'],
-                         'squeeze')
+    check_type_and_dtype(
+        input, 'input', Variable,
+        ['float16', 'float32', 'float64', 'int8', 'int32', 'int64'], 'squeeze')
     check_type(axes, 'axes', list, 'squeeze')
     out = helper.create_variable_for_type_inference(dtype=input.dtype)
     x_shape = helper.create_variable_for_type_inference(dtype=input.dtype)

python/paddle/fluid/layers/tensor.py

@@ -448,7 +448,7 @@ def assign(input, output=None):
 
     Parameters:
         input (Variable|numpy.ndarray): A tensor or numpy ndarray, its data type supports
-            float32, float64, int32 and int64.
+            float16, float32, float64, int32 and int64.
         output (Variable, optional): A tensor. If :attr:`output` is None, a new tensor will
             be created as :attr:`output`. Default: None.
 
@@ -469,9 +469,10 @@ def assign(input, output=None):
     helper = LayerHelper('assign', **locals())
     check_type(input, 'input', (Variable, numpy.ndarray), 'assign')
     if isinstance(input, Variable):
-        check_dtype(input.dtype, 'input',
-                    ['float32', 'float64', 'int32', 'int64', 'bool'], 'assign',
-                    '(When the type of input in assign is Variable.)')
+        check_dtype(
+            input.dtype, 'input',
+            ['float16', 'float32', 'float64', 'int32', 'int64', 'bool'],
+            'assign', '(When the type of input in assign is Variable.)')
         if output is None:
             output = helper.create_variable_for_type_inference(
                 dtype=input.dtype)