!49450 [MSLITE] Fix document of topk operator.

Merge pull request !49450 from wangshaocong/bugfix

docs/api/api_python/ops/mindspore.ops.TopK.rst

@@ -21,7 +21,7 @@
         - **sorted** (bool, 可选) - 如果为True，则获取的元素将按值降序排序。如果为False，则不对获取的元素进行排序。默认值：True。
 
     输入：
-        - **input_x** (Tensor) - 需计算的输入，数据类型必须为float16、float32或int32。
+        - **input_x** (Tensor) - 需计算的输入，CPU推理数据类型必须为float16、float32或int32；GPU推理数据类型必须为float16或float32。
         - **k** (int) - 指定计算最大元素的数量，必须为常量。
 
     输出：

mindspore/ccsrc/plugin/device/gpu/kernel/arrays/topk_gpu_kernel.cc

@@ -36,14 +36,6 @@ MS_REG_GPU_KERNEL_TWO(TopK,
                         .AddOutputAttr(kNumberTypeInt32),
                       TopKGpuKernelMod, half, int)
 
-MS_REG_GPU_KERNEL_TWO(TopK,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeInt32),
-                      TopKGpuKernelMod, int, int)
-
 MS_REG_GPU_KERNEL_TWO(TopK,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeFloat32)
@@ -59,13 +51,5 @@ MS_REG_GPU_KERNEL_TWO(TopK,
                         .AddOutputAttr(kNumberTypeFloat16)
                         .AddOutputAttr(kNumberTypeInt32),
                       TopKGpuKernelMod, half, int)
-
-MS_REG_GPU_KERNEL_TWO(TopK,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeInt32),
-                      TopKGpuKernelMod, int, int)
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_ops/topk_impl.cu

@@ -226,5 +226,3 @@ template CUDA_LIB_EXPORT void FastTopK(const int outer_size, const int inner_siz
                                        half *output, int *output_index, const half init_K, cudaStream_t stream);
 template CUDA_LIB_EXPORT void FastTopK(const int outer_size, const int inner_size, const float *input, int k_cut,
                                        float *output, int *output_index, const float init_K, cudaStream_t stream);
-template CUDA_LIB_EXPORT void FastTopK(const int outer_size, const int inner_size, const int *input, int k_cut,
-                                       int *output, int *output_index, const int init_K, cudaStream_t stream);

mindspore/python/mindspore/ops/operations/array_ops.py

@@ -7950,7 +7950,8 @@ class TopK(Primitive):
             If False, the obtained elements will not be sorted. Default: True.
 
     Inputs:
-        - **input_x** (Tensor) - Input to be computed, data type must be float16, float32 or int32.
+        - **input_x** (Tensor) - Input to be computed, data type must be float16, float32 or int32 on CPU,
+          and float16 or float32 on GPU.
         - **k** (int) - The number of top elements to be computed along the last dimension, constant input is needed.
 
     Outputs: