!34597 add intopk vmap, functional/tensor api

Merge pull request !34597 from 范吉斌/intopk_vmap

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

@@ -215,6 +215,7 @@ Reduction算子
     mindspore.ops.equal
     mindspore.ops.ge
     mindspore.ops.gt
+    mindspore.ops.intopk
     mindspore.ops.isfinite
     mindspore.ops.isnan
     mindspore.ops.le

docs/api/api_python/mindspore/mindspore.Tensor.rst

@@ -593,6 +593,21 @@ mindspore.Tensor
         - **TypeError** - `indices` 不是int或tuple。
         - **TypeError** - `indices` 是元组，但是其中的元素不是int。
 
+    .. py:method:: intopk(x, k)
+
+        判断目标标签是否在前 `k` 个预测中。
+
+        更多细节参考 :func:`mindspore.ops.intopk`。
+
+        **参数：**
+
+        **x** (Tensor) - 1维Tensor，样本的标签。数据类型为int32。 `x` 的大小必须与该Tensor第一维度的大小相同。 `x` 取值不可为负且必须小于或等于该Tensor第二维度的大小。
+        **k** (int32) - 指定在最后一维上参与比较的top元素的数量。
+
+        **返回：**
+
+        1维的bool类型Tensor，与 `x` shape相同。对于 `x` 中的样本标签 `i`，如果它在该Tensor的前 `k` 个预测值中，则输出值为True，否则为False。
+
     .. py:method:: item(index=None)
 
         获取Tensor中指定索引的元素。
@@ -1508,3 +1523,4 @@ mindspore.Tensor
 
         Tensor，具有与入参 `shape` 相同的维度。
     
+

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

@@ -7,19 +7,19 @@ mindspore.ops.InTopK
 
     **参数：**
     
-    - **k** (int) - 指定 `k` 的值。
+    - **k** (int) - 指定在最后一维上参与比较的top元素的数量。
 
     **输入：**
-    
+
     - **x1** (Tensor) - 2维Tensor，对样本的预测。数据类型支持float16或float32。
     - **x2** (Tensor) - 1维Tensor，样本的标签。数据类型为int32。 `x2` 的大小必须与 `x1` 第一维度的大小相同。 `x2` 取值不可为负且必须小于或等于 `x1` 第二维度的大小。
 
     **输出：**
-    
-    1维的bool类型Tensor，与 `x2` shape相同。对于样本 `i` 在 `x2` 中的标签, 如果其在 `x1` 对样本 `i` 的前 `k` 个预测中，则输出值为True，否则为False。
+
+    1维的bool类型Tensor，与 `x2` shape相同。对于 `x2` 中的样本标签 `i`，如果它在 `x1` 的前 `k` 个预测值中，则输出值为True，否则为False。
 
     **异常：**
-    
+
     - **TypeError** - `k` 不是int类型。
     - **TypeError** - `x1` 或 `x2` 不是Tensor。
     - **TypeError** - `x1` 的数据类型非float16或float32。

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

@@ -0,0 +1,22 @@
+mindspore.ops.intopk
+====================
+
+.. py:function:: mindspore.ops.intopk(x1, x2, k)
+
+    判断目标标签是否在前 `k` 个预测中。
+
+    **参数：**
+
+    **x1** (Tensor) - 2维Tensor，对样本的预测。数据类型支持float16或float32。
+    **x2** (Tensor) - 1维Tensor，样本的标签。数据类型为int32。 `x2` 的大小必须与 `x1` 第一维度的大小相同。 `x2` 取值不可为负且必须小于或等于 `x1` 第二维度的大小。
+    **k** (int32) - 指定在最后一维上参与比较的top元素的数量。
+
+    **输出：**
+
+    1维的bool类型Tensor，与 `x2` shape相同。对于 `x2` 中的样本标签 `i`，如果它在 `x1` 的前 `k` 个预测值中，则输出值为True，否则为False。
+
+    **异常：**
+
+    - **TypeError** - `k` 不是int类型。
+    - **TypeError** - `x1` 或 `x2` 不是Tensor。
+    - **TypeError** - `x1` 的数据类型非float16或float32。

docs/api/api_python_en/mindspore.ops.functional.rst

@@ -214,6 +214,7 @@ Comparison operators
     mindspore.ops.equal
     mindspore.ops.ge
     mindspore.ops.gt
+    mindspore.ops.intopk
     mindspore.ops.isfinite
     mindspore.ops.isnan
     mindspore.ops.le
@@ -302,6 +303,7 @@ Array Operation
     mindspore.ops.gather
     mindspore.ops.gather_d
     mindspore.ops.gather_nd
+    mindspore.ops.intopk
     mindspore.ops.masked_select
     mindspore.ops.matrix_band_part
     mindspore.ops.matrix_diag

mindspore/ccsrc/pipeline/jit/resource.cc

@@ -246,6 +246,7 @@ BuiltInTypeMap &GetMethodMap() {
        {"cdist", std::string("cdist")},                                    // P.cdist
        {"hardshrink", std::string("hardshrink")},                          // P.hshrink
        {"one_hot", std::string("one_hot")},                                // P.OneHot
+       {"intopk", std::string("intopk")},                                  // P.InTopK
        {"gather_nd", std::string("gather_nd")},                            // P.GatherNd()
        {"unique_consecutive", std::string("unique_consecutive")},          // UniqueConsecutive()
        {"diag", std::string("diag")},                                      // P.Diag()

mindspore/python/mindspore/_extends/parse/standard_method.py

@@ -1642,6 +1642,14 @@ def masked_fill(x, mask, value):
     return C.array_ops.masked_fill(x, mask, value)
 
 
+def intopk(x1, x2, k):
+    """
+    Determines whether the targets are in the top `k` predictions.
+    """
+    check_is_int(k, 'k')
+    return F.intopk(x1, x2, k)
+
+
 def narrow(x, axis, start, length):
     """
     Returns a narrowed tensor from input tensor.

mindspore/python/mindspore/common/tensor.py

@@ -3800,6 +3800,29 @@ class Tensor(Tensor_):
         self._init_check()
         return tensor_operator_registry.get('diag')()(self)
 
+    def intopk(self, x, k):
+        r"""
+        Determines whether the targets are in the top `k` predictions.
+
+        Refer to :func:`mindspore.ops.intopk` for more detail.
+
+        Supported Platforms:
+            ``Ascend`` ``GPU`` ``CPU``
+
+        Examples:
+            >>> x1 = Tensor(np.array([[1, 8, 5, 2, 7], [4, 9, 1, 3, 5]]), mindspore.float32)
+            >>> x2 = Tensor(np.array([1, 3]), mindspore.int32)
+            >>> output = x1.intopk(x2, 3)
+            >>> print(output)
+            [ True  False]
+        """
+        if not isinstance(x, Tensor):
+            raise TypeError("For 'Tensor.intopk', the type of the argument 'x' must be Tensor, but "
+                            "got {}.".format(type(x)))
+        validator.check_type_name('x', x.dtype, [mstype.float16, mstype.float32], "Tensor")
+        validator.check_value_type("k", k, [int], "Tensor")
+        return tensor_operator_registry.get('intopk')(self, x, k)
+
 
 class RowTensor(RowTensor_):
     """

mindspore/python/mindspore/ops/_vmap/vmap_nn_ops.py

@@ -123,6 +123,30 @@ def get_dropout_nd_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register(P.InTopK)
+def get_in_top_k_vmap_rule(prim, axis_size):
+    """VmapRule for `InTopK`."""
+
+    def vmap_rule(x1_bdim, x2_bdim):
+        is_all_none, result = vmap_general_preprocess(prim, x1_bdim, x2_bdim)
+        if is_all_none:
+            return result
+
+        x1, x1_dim = x1_bdim
+        x2, x2_dim = x2_bdim
+        x1 = _bdim_at_front(x1, x1_dim, axis_size)
+        x2 = _bdim_at_front(x2, x2_dim, axis_size)
+        x1_shape = F.shape(x1)
+        x2_shape = F.shape(x2)
+        x1 = F.reshape(x1, (-1, x1_shape[-1]))
+        x2 = F.reshape(x2, (-1,))
+        output = prim(x1, x2)
+        output = F.reshape(output, x2_shape)
+        return (output, 0)
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register(G.FastGeLUGrad)
 @vmap_rules_getters.register(G.HShrinkGrad)
 def get_fast_gelu_grad_vmap_rule(prim, axis_size):

mindspore/python/mindspore/ops/function/__init__.py

@@ -199,6 +199,7 @@ from .nn_func import (
     deformable_conv2d,
     fast_gelu,
     hardshrink,
+    intopk,
     softsign,
     pdist,
     nll_loss,

mindspore/python/mindspore/ops/function/nn_func.py

@@ -461,10 +461,45 @@ def _nll_loss(inputs, target, target_dim=-1, weight=None, ignore_index=None, red
     return loss
 
 
+def intopk(x1, x2, k):
+    r"""
+    Determines whether the targets are in the top `k` predictions.
+
+    Args:
+        x1 (Tensor): A 2D Tensor defines the predictions of a batch of samples with float16 or float32
+          data type.
+        x2 (Tensor): A 1D Tensor defines the labels of a batch of samples with int32 data type. The size of `x2`
+          must be equal to the first dimension of `x1`. The values of `x2` can not be negative and
+          must be equal to or less than index of x1's second dimension.
+        k (int): Specifies the number of top elements to be used for computing precision along the last dimension.
+
+    Returns:
+        Tensor has 1 dimension of type bool and the same shape with `x2`. For labeling sample `i` in `x2`,
+        if the label in the first `k` predictions for sample `i` is in `x1`, then the value is True, otherwise False.
+
+    Raises:
+        TypeError: If `k` is not an int.
+        TypeError: If `x1` or `x2` is not a Tensor.
+        TypeError: If dtype of `x1` is neither float16 nor float32.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> x1 = Tensor(np.array([[1, 8, 5, 2, 7], [4, 9, 1, 3, 5]]), mindspore.float32)
+        >>> x2 = Tensor(np.array([1, 3]), mindspore.int32)
+        >>> output = ops.intopk(x1, x2, 3)
+        >>> print(output)
+        [ True  False]
+    """
+    return P.InTopK(k)(x1, x2)
+
+
 __all__ = [
     'deformable_conv2d',
     'fast_gelu',
     'hardshrink',
+    'intopk',
     'softsign',
     'pdist',
     'cross_entropy',

mindspore/python/mindspore/ops/functional.py

@@ -979,6 +979,7 @@ tensor_operator_registry.register('stack', P.Stack)
 tensor_operator_registry.register('log', log)
 tensor_operator_registry.register('lerp', lerp)
 tensor_operator_registry.register('floor', floor)
+tensor_operator_registry.register('intopk', intopk)
 # support sparse tensor operators
 tensor_operator_registry.register('csr_mul', csr_mul)
 tensor_operator_registry.register('csr2coo', csr2coo)

mindspore/python/mindspore/ops/operations/nn_ops.py

@@ -7680,27 +7680,10 @@ class InTopK(PrimitiveWithInfer):
     r"""
     Determines whether the targets are in the top `k` predictions.
 
-    Args:
-        k (int): Specifies the number of top elements to be used for computing precision.
-
-    Inputs:
-        - **x1** (Tensor) - A 2D Tensor defines the predictions of a batch of samples with float16 or float32
-          data type.
-        - **x2** (Tensor) - A 1D Tensor defines the labels of a batch of samples with int32 data type. The size of x2
-          must be equal to x1's first dimension. The values of `x2` can not be negative and
-          must be equal to or less than index of x1's second dimension.
-
-    Outputs:
-        Tensor has 1 dimension of type bool and the same shape with `x2`. For labeling sample `i` in `x2`,
-        if the label in the first `k` predictions for sample `i` is in `x1`, then the value is True, otherwise False.
-
-    Raises:
-        TypeError: If `k` is not an int.
-        TypeError: If `x1` or `x2` is not a Tensor.
-        TypeError: If dtype of `x1` is neither float16 nor float32.
+    Refer to :func:`mindspore.ops.tanh` for more detail.
 
     Supported Platforms:
-        ``Ascend`` ``GPU``
+        ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
         >>> x1 = Tensor(np.array([[1, 8, 5, 2, 7], [4, 9, 1, 3, 5]]), mindspore.float32)