未验证 提交 b8eb52f6 编写于 作者: W wangchaochaohu 提交者: GitHub

fix the fp16 support of assgin and squeeze Op test=develop (#24939)

上级 dc7d34ef
......@@ -7479,7 +7479,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)`.
......@@ -7499,9 +7499,9 @@ def squeeze(input, axes, name=None):
"""
helper = LayerHelper("squeeze", **locals())
check_variable_and_dtype(input, 'input',
['float32', 'float64', 'int8', 'int32', 'int64'],
'squeeze')
check_variable_and_dtype(
input, 'input',
['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)
......
......@@ -566,7 +566,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.
......@@ -587,9 +587,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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册