!4264 Modified some API about supported data type for ops.

Merge pull request !4264 from liuxiao93/check-fp16-fp32

mindspore/ops/operations/array_ops.py

@@ -1423,7 +1423,9 @@ class UnsortedSegmentMin(PrimitiveWithInfer):
 
     Inputs:
         - **input_x** (Tensor) - The shape is :math:`(x_1, x_2, ..., x_R)`.
+          The data type should be float16, float32 or int32.
         - **segment_ids** (Tensor) - A `1-D` tensor whose shape is :math:`(x_1)`, the value should be >= 0.
+          The data type must be int32.
         - **num_segments** (int) - The value spcifies the number of distinct `segment_ids`.
 
     Outputs:
@@ -2410,7 +2412,7 @@ class GatherNd(PrimitiveWithInfer):
 
     Inputs:
         - **input_x** (Tensor) - The target tensor to gather values.
-        - **indices** (Tensor) - The index tensor.
+        - **indices** (Tensor) - The index tensor, with int data type.
 
     Outputs:
         Tensor, has the same type as `input_x` and the shape is indices_shape[:-1] + x_shape[indices_shape[-1]:].
@@ -2807,7 +2809,7 @@ class ScatterNonAliasingAdd(_ScatterNdOp):
     This operation outputs the `input_x` after the update is done, which makes it convenient to use the updated value.
 
     Inputs:
-        - **input_x** (Parameter) - The target parameter.
+        - **input_x** (Parameter) - The target parameter. The data type should be float16, float32 or int32.
         - **indices** (Tensor) - The index to do add operation whose data type should be mindspore.int32.
         - **updates** (Tensor) - The tensor doing the add operation with `input_x`,
           the data type is same as `input_x`, the shape is `indices_shape[:-1] + x_shape[indices_shape[-1]:]`.

mindspore/ops/operations/math_ops.py

@@ -943,9 +943,9 @@ class InplaceAdd(PrimitiveWithInfer):
             to add with v. It is a int or tuple, whose value is in [0, the first dimension size of x).
 
     Inputs:
-        - **input_x** (Tensor) - The first input is a tensor whose data type is number.
+        - **input_x** (Tensor) - The first input is a tensor whose data type is float16, float32 or int32.
         - **input_v** (Tensor) - The second input is a tensor who has the same dimension sizes as x except
-          the first dimension, which must be the same as indices's size.
+          the first dimension, which must be the same as indices's size. It has the same data type with `input_x`.
 
     Outputs:
         Tensor, has the same shape and dtype as input.
@@ -1001,9 +1001,9 @@ class InplaceSub(PrimitiveWithInfer):
             to sub with v. It is a int or tuple, whose value is in [0, the first dimension size of x).
 
     Inputs:
-        - **input_x** (Tensor) - The first input is a tensor whose data type is number.
+        - **input_x** (Tensor) - The first input is a tensor whose data type is float16, float32 or int32.
         - **input_v** (Tensor) - The second input is a tensor who has the same dimension sizes as x except
-          the first dimension, which must be the same as indices's size.
+          the first dimension, which must be the same as indices's size. It has the same data type with `input_x`.
 
     Outputs:
         Tensor, has the same shape and dtype as input.
@@ -1403,7 +1403,7 @@ class Expm1(PrimitiveWithInfer):
     Returns exponential then minus 1 of a tensor element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor.
+        - **input_x** (Tensor) - The input tensor. With float16 or float32 data type.
 
     Outputs:
         Tensor, has the same shape as the `input_x`.
@@ -1425,6 +1425,7 @@ class Expm1(PrimitiveWithInfer):
 
     def infer_dtype(self, x_type):
         validator.check_subclass("x", x_type, mstype.tensor, self.name)
+        validator.check_tensor_type_same({"x": x_type}, [mstype.float16, mstype.float32], self.name)
         return x_type
 
 
@@ -1515,7 +1516,7 @@ class Log1p(PrimitiveWithInfer):
     Returns the natural logarithm of one plus the input tensor element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor.
+        - **input_x** (Tensor) - The input tensor. With float16 or float32 data type.
 
     Outputs:
         Tensor, has the same shape as the `input_x`.
@@ -1536,6 +1537,7 @@ class Log1p(PrimitiveWithInfer):
 
     def infer_dtype(self, x):
         validator.check_subclass("x", x, mstype.tensor, self.name)
+        validator.check_tensor_type_same({"x": x}, [mstype.float16, mstype.float32], self.name)
         return x
 
 
@@ -1544,7 +1546,7 @@ class Erf(PrimitiveWithInfer):
     Computes the Gauss error function of `input_x` element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor.
+        - **input_x** (Tensor) - The input tensor. The data type must be float16 or float32.
 
     Outputs:
         Tensor, has the same shape and dtype as the `input_x`.
@@ -1574,7 +1576,7 @@ class Erfc(PrimitiveWithInfer):
     Computes the complementary error function of `input_x` element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor.
+        - **input_x** (Tensor) - The input tensor. The data type mast be float16 or float32.
 
     Outputs:
         Tensor, has the same shape and dtype as the `input_x`.
@@ -1674,6 +1676,7 @@ class Maximum(_MathBinaryOp):
             return Tensor(out)
         return None
 
+
 class RealDiv(_MathBinaryOp):
     """
     Divide the first input tensor by the second input tensor in floating-point type element-wise.
@@ -1923,7 +1926,7 @@ class Floor(PrimitiveWithInfer):
     Round a tensor down to the closest integer element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor. Its element data type must be float.
+        - **input_x** (Tensor) - The input tensor. It's element data type must be float.
 
     Outputs:
         Tensor, has the same shape as `input_x`.
@@ -1981,7 +1984,7 @@ class Ceil(PrimitiveWithInfer):
     Round a tensor up to the closest integer element-wise.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor. Its element data type must be float.
+        - **input_x** (Tensor) - The input tensor. It's element data type must be float16 or float32.
 
     Outputs:
         Tensor, has the same shape as `input_x`.
@@ -2001,7 +2004,7 @@ class Ceil(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_dtype):
-        validator.check_tensor_type_same({"x": x_dtype}, mstype.float_type, self.name)
+        validator.check_tensor_type_same({"x": x_dtype}, [mstype.float16, mstype.float32], self.name)
         return x_dtype
 
 
@@ -2666,7 +2669,7 @@ class FloatStatus(PrimitiveWithInfer):
     Determine if the elements contains nan, inf or -inf. `0` for normal, `1` for overflow.
 
     Inputs:
-        - **input_x** (Tensor) - The input tensor.
+        - **input_x** (Tensor) - The input tensor. The data type must be float16 or float32.
 
     Outputs:
         Tensor, has the shape of `(1,)`, and has the same dtype of input `mindspore.dtype.float32` or
@@ -2731,6 +2734,7 @@ class NPUGetFloatStatus(PrimitiveWithInfer):
 
     Inputs:
         - **input_x** (Tensor) - The output tensor of `NPUAllocFloatStatus`.
+          The data type must be float16 or float32.
 
     Outputs:
         Tensor, has the same shape as `input_x`. All the elements in the tensor will be zero.
@@ -2755,7 +2759,7 @@ class NPUGetFloatStatus(PrimitiveWithInfer):
         return [8]
 
     def infer_dtype(self, x_dtype):
-        validator.check_tensor_type_same({'x': x_dtype}, [mstype.float32], self.name)
+        validator.check_tensor_type_same({'x': x_dtype}, [mstype.float16, mstype.float32], self.name)
         return mstype.float32
 
 
@@ -2771,6 +2775,7 @@ class NPUClearFloatStatus(PrimitiveWithInfer):
 
     Inputs:
         - **input_x** (Tensor) - The output tensor of `NPUAllocFloatStatus`.
+          The data type must be float16 or float32.
 
     Outputs:
         Tensor, has the same shape as `input_x`. All the elements in the tensor will be zero.
@@ -2797,7 +2802,7 @@ class NPUClearFloatStatus(PrimitiveWithInfer):
         return [8]
 
     def infer_dtype(self, x_dtype):
-        validator.check_tensor_type_same({'x': x_dtype}, [mstype.float32], self.name)
+        validator.check_tensor_type_same({'x': x_dtype}, [mstype.float16, mstype.float32], self.name)
         return mstype.float32
 
 
@@ -2932,6 +2937,7 @@ class NMSWithMask(PrimitiveWithInfer):
           `N` is the number of input bounding boxes. Every bounding box
           contains 5 values, the first 4 values are the coordinates of bounding
           box, and the last value is the score of this bounding box.
+          The data type must be float16 or float32.
 
     Outputs:
         tuple[Tensor], tuple of three tensors, they are selected_boxes, selected_idx and selected_mask.
@@ -3186,12 +3192,13 @@ class Atan2(_MathBinaryOp):
          [[0. 0.7853982]]
     """
 
+
 class SquareSumAll(PrimitiveWithInfer):
     """
     Returns square sum all of a tensor element-wise
 
     Inputs:
-        - **input_x1** (Tensor) - The input tensor.
+        - **input_x1** (Tensor) - The input tensor. The data type must be float16 or float32.
         - **input_x2** (Tensor) - The input tensor same type and shape as the `input_x1`.
 
     Note:
@@ -3227,7 +3234,7 @@ class BitwiseAnd(_BitwiseBinaryOp):
     Returns bitwise `and` of two tensors element-wise.
 
     Inputs:
-        - **input_x1** (Tensor) - The input tensor with int or uint type.
+        - **input_x1** (Tensor) - The input tensor with int16 or uint16 data type.
         - **input_x2** (Tensor) - The input tensor with same type as the `input_x1`.
 
     Outputs:
@@ -3247,7 +3254,7 @@ class BitwiseOr(_BitwiseBinaryOp):
     Returns bitwise `or` of two tensors element-wise.
 
     Inputs:
-        - **input_x1** (Tensor) - The input tensor with int or uint type.
+        - **input_x1** (Tensor) - The input tensor with int16 or uint16 data type.
         - **input_x2** (Tensor) - The input tensor with same type as the `input_x1`.
 
     Outputs:
@@ -3267,7 +3274,7 @@ class BitwiseXor(_BitwiseBinaryOp):
     Returns bitwise `xor` of two tensors element-wise.
 
     Inputs:
-        - **input_x1** (Tensor) - The input tensor with int or uint type.
+        - **input_x1** (Tensor) - The input tensor with int16 or uint16 data type.
         - **input_x2** (Tensor) - The input tensor with same type as the `input_x1`.
 
     Outputs:
@@ -3405,7 +3412,7 @@ class Eps(PrimitiveWithInfer):
     Creates a tensor filled with `input_x` dtype minimum val.
 
     Inputs:
-        - **input_x** (Tensor) - Input tensor.
+        - **input_x** (Tensor) - Input tensor. The data type must be float16 or float32.
 
     Outputs:
         Tensor, has the same type and shape as `input_x`, but filled with `input_x` dtype minimum val.

mindspore/ops/operations/other_ops.py

@@ -181,8 +181,9 @@ class CheckValid(PrimitiveWithInfer):
     Check whether the bounding box cross data and data border.
 
     Inputs:
-        - **bboxes** (Tensor) - Bounding boxes tensor with shape (N, 4).
+        - **bboxes** (Tensor) - Bounding boxes tensor with shape (N, 4). Data type should be float16 or float32.
         - **img_metas** (Tensor) - Raw image size information, format (height, width, ratio).
+          Data type should be float16 or float32.
 
     Outputs:
         Tensor, the valided tensor.
@@ -220,6 +221,9 @@ class CheckValid(PrimitiveWithInfer):
         return bboxes_shape[:-1]
 
     def infer_dtype(self, bboxes_type, metas_type):
+        valid_type = [mstype.float32, mstype.float16]
+        validator.check_tensor_type_same({"bboxes_type": bboxes_type}, valid_type, self.name)
+        validator.check_tensor_type_same({"metas_type": metas_type}, valid_type, self.name)
         return mstype.bool_
 
 
@@ -242,12 +246,12 @@ class IOU(PrimitiveWithInfer):
 
     Inputs:
         - **anchor_boxes** (Tensor) - Anchor boxes, tensor of shape (N, 4). "N" indicates the number of anchor boxes,
-          and the value "4" refers to "x0", "x1", "y0", and "y1". Data type must be float16.
+          and the value "4" refers to "x0", "x1", "y0", and "y1". Data type must be float16 or float32.
         - **gt_boxes** (Tensor) - Ground truth boxes, tensor of shape (M, 4). "M" indicates the number of ground
-          truth boxes, and the value "4" refers to "x0", "x1", "y0", and "y1". Data type must be float16.
+          truth boxes, and the value "4" refers to "x0", "x1", "y0", and "y1". Data type must be float16 or float32.
 
     Outputs:
-        Tensor, the 'iou' values, tensor of shape (M, N), with data type float16.
+        Tensor, the 'iou' values, tensor of shape (M, N), with the same data type as `anchor_boxes`.
 
     Raises:
         KeyError: When `mode` is not 'iou' or 'iof'.
@@ -274,6 +278,9 @@ class IOU(PrimitiveWithInfer):
         return iou
 
     def infer_dtype(self, anchor_boxes, gt_boxes):
+        valid_type = [mstype.float32, mstype.float16]
+        validator.check_tensor_type_same({"anchor_boxes": anchor_boxes}, valid_type, self.name)
+        validator.check_tensor_type_same({"gt_boxes": gt_boxes}, valid_type, self.name)
         return anchor_boxes