Argmax
=================


沿指定轴查找最大 `topk` 个值的索引。当 `topk=1` 时，该算子等价于 `ArgMax`。

.. math::

    Y_i = \underset{k}{\operatorname{argmax}} (X_{slice_i})

其中 :math:`X_{slice_i}` 是输入张量中沿指定轴的一个切片，函数返回该切片中最大值的索引 :math:`k`。

输入：
    - **input** - 输入数据地址。
    - **output** - 输出索引的数据地址，数据类型通常为int32。
    - **output_value** - (可选) 输出值的数据地址。
    - **in_shape** - 输入张量的维度信息数组。
    - **in_strides** - 输入张量的步长信息数组。
    - **out_strides** - 输出张量的步长信息数组。
    - **arg_elements** - 用于存放候选值的临时工作空间地址。
    - **index** - 用于存放候选索引的临时工作空间地址。
    - **topk** - 需要查找的最大值的数量。设置为1以执行ArgMax操作。
    - **out_value** - 是否返回数值的标志。若为非0，则 `output_value` 必须提供有效地址。
    - **input_shape_size** - 输入张量的维度数 (即 `in_shape` 数组的长度)。
    - **axis** - 执行查找操作的轴。
    - **core_mask** - 核掩码。

输出：
    - **output** - 存储索引的输出张量。
    - **output_value** - 如果 `return_values` 为 true，则此处存储找到的值。

支持平台：
    ``FT78NE``
    ``MT7004``

.. note::
    - FT78NE 支持fp32
    - MT7004 支持fp16, fp32


**共享存储版本:**

.. c:function:: void fp_argmax_s(float *input, void *output, float *output_value, int32_t *in_shape, int *in_strides, int *out_strides, float *arg_elements, int* index, int topk, int out_value, int input_shape_size, int axis, int core_mask)
.. c:function:: void hp_argmax_s(half *input, void *output, half *output_value, int32_t *in_shape, int *in_strides, int *out_strides, half *arg_elements, int* index, int topk, int out_value, int input_shape_size, int axis, int core_mask)

**C调用示例：**

.. code-block:: c
    :linenos:
    :emphasize-lines: 21

    //FT78NE示例
    #include <stdio.h>
    #include <argmax.h>
    int main(int argc, char* argv[]) {
        float *input = (float *)0xA0000000;          // input在DDR空间
        int *output = (int *)0xB0000000;             // output indices
        float *output_value = (float *)0xC0000000;     // output values
        float *arg_elements = (float *)0xD0000000;   // temp workspace 1
        int *index = (int *)0xE0000000;               // temp workspace 2

        int in_shape[] = {2, 3, 4};                  // input shape: (2, 3, 4)
        int in_strides[] = {12, 4, 1};               // input strides for contiguous layout
        int out_strides[] = {4, 1};                  // output strides, shape is (2, 4)
        int input_shape_size = 3;
        
        int axis = 1;                                // 沿第1轴操作
        int topk = 1;                                // ArgMax
        int out_value = 1;                           // 同时返回值
        int core_mask = 0xff;
        
        fp_argmax_s(input, output, output_value, in_shape, in_strides, out_strides, arg_elements, index, topk, out_value, input_shape_size, axis, core_mask);
        return 0;
    }


**私有存储版本:**

.. c:function:: void fp_argmax_p(float *input, void *output, float *output_value, int32_t *in_shape, int *in_strides, int *out_strides, float *arg_elements, int* index, int topk, int out_value, int input_shape_size, int axis)
.. c:function:: void hp_argmax_p(half *input, void *output, half *output_value, int32_t *in_shape, int *in_strides, int *out_strides, half *arg_elements, int* index, int topk, int out_value, int input_shape_size, int axis)

   
**C调用示例：**

.. code-block:: c
    :linenos:
    :emphasize-lines: 20-21

    //FT78NE示例
    #include <stdio.h>
    #include <argmax.h>
    int main(int argc, char* argv[]) {
        float *input = (float *)0x10001000;          // input在DDR空间
        int *output = (int *)0x10002000;             // output indices
        float *output_value = (float *)0x10003000;     // output values
        float *arg_elements = (float *)0x10004000;   // temp workspace 1
        int *index = (int *)0x10005000;               // temp workspace 2

        int in_shape[] = {2, 3, 4};                  // input shape: (2, 3, 4)
        int in_strides[] = {12, 4, 1};               // input strides for contiguous layout
        int out_strides[] = {4, 1};                  // output strides, shape is (2, 4)
        int input_shape_size = 3;
        
        int axis = 1;                                // 沿第1轴操作
        int topk = 1;                                // ArgMax
        int out_value = 1;                           // 同时返回值
        
        fp_argmax_p(input, output, output_value, in_shape, in_strides, out_strides, 
                    arg_elements, index, topk, out_value, input_shape_size, axis);
        return 0;
    }