ArgMax 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介返回张量在指定维度上的最大值的索引。主要应用场景分类任务中获取预测类别取 logits 最大值对应的类别索引Top-1 准确率计算贪心解码Greedy Decoding中选择概率最大的 token算子特征难度等级L2SortSelect单输入单输出沿指定维度进行归约操作输入支持 0-8 维输出维度比输入少一维沿指定 dimension 维度归约2. 算子定义数学公式$$ y \arg\max_{axisdimension}(x) $$即返回输入张量 $x$ 在指定维度dimension上最大值所在的索引位置。3. 接口规范算子原型cann_bench.arg_max(Tensor x, int dimension) - Tensor y输入参数说明参数类型默认值描述xTensor必选输入张量dimensionint64必选计算 argmax 的维度输出参数Shapedtype描述y输入 x 去掉 dimension 维后的 shapeint64输出张量最大值的索引数据类型输入 dtype输出 dtypefloat16int64float32int64bfloat16int64int32int64int64int64规则与约束输入支持 0-8 维张量dimension支持负数索引如 -1 表示最后一维输出 dtype 固定为 int64输出 shape 为输入 shape 去掉dimension维度后的结果当指定维度上存在多个相同的最大值时返回第一个最小索引出现的位置4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch ArgMax 算子 Torch Golden 参考实现 返回张量在指定维度上的最大值的索引 公式: y argmax(x, axisdimension) def arg_max( x: torch.Tensor, dimension: int ) - torch.Tensor: 返回张量在指定维度上的最大值的索引 公式: y argmax(x, axisdimension) Args: x: 输入张量 dimension: 计算 argmax 的维度 Returns: 输出张量最大值的索引int64 y torch.argmax(x, dimdimension) return y6. 额外信息算子调用示例import torch import cann_bench x torch.randn(1024, 1024, dtypetorch.float32, devicenpu) y cann_bench.arg_max(x, dimension-1) # 沿最后一维取 argmax输出 shape [1024] x torch.randn(2, 8, 256, 256, dtypetorch.float16, devicenpu) y cann_bench.arg_max(x, dimension2) # 沿第 2 维取 argmax输出 shape [2, 8, 256]【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考