add support data type to ArgMax and fix some error about Floor operator

This commit is contained in:
dinglinhe 2021-07-30 09:27:19 +08:00
parent 7187887b63
commit bf5a814672
2 changed files with 8 additions and 3 deletions

View File

@ -1654,7 +1654,12 @@ class Argmax(PrimitiveWithInfer):
Default: `mindspore.dtype.int32`.
Inputs:
- **input_x** (Tensor) - Input tensor.
- **input_x** (Tensor) - Input tensor. :math:`(N,*)` where :math:`*` means, any number of additional dimensions.
Support data type list as follows:
- Ascend: Float16, Float32.
- GPU: Float16, Float32.
- CPU: Float16, Float32, Float64.
Outputs:
Tensor, indices of the max value of input tensor across the axis.

View File

@ -2634,7 +2634,7 @@ class FloorDiv(_MathBinaryOp):
out_{i} = \\text{floor}( \\frac{x_i}{y_i})
where the :math:`floor` indicates the operator that converts the input data into the floor data type.
where the :math:`floor` indicates the Floor operator, for more details, please refer to the Floor operator.
Inputs:
- **x** (Union[Tensor, Number, bool]) - The first input is a number or
@ -2841,7 +2841,7 @@ class FloorMod(_MathBinaryOp):
out_{i} =\text{floor}(x_{i} // y_{i})
where the :math:`floor` indicates the operator that converts the input data into the floor data type.
where the :math:`floor` indicates the Floor operator, for more details, please refer to the Floor operator.
Inputs:
- **x** (Union[Tensor, Number, bool]) - The first input is a number or