ArgMaxV2

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	×
Atlas 推理系列产品	√
Atlas 训练系列产品	√

功能说明

返回张量在指定维度（dim）上的最大值的索引，并保存到out张量中。

函数原型

每个算子分为两段式接口，必须先调用“aclnnArgMaxGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnArgMax”接口执行计算。

aclnnStatus aclnnArgMaxGetWorkspaceSize(
  const aclTensor *self, 
  int64_t          dim, 
  bool             keepdim, 
  aclTensor       *out, 
  uint64_t        *workspaceSize, 
  aclOpExecutor  **executor)

aclnnStatus aclnnArgMax(
  void          *workspace, 
  uint64_t       workspaceSize, 
  aclOpExecutor *executor, 
  aclrtStream    stream)

aclnnArgMaxGetWorkspaceSize

• 参数说明：

  参数名	输入/输出	描述	使用说明	数据类型	数据格式	维度(shape)	非连续Tensor
  self	输入	输入Tensor。	-	FLOAT、FLOAT16、BFLOAT16、DOUBLE、INT8、INT16、INT32、INT64、UINT8	ND	-	√
  dim	输入	指定了要进行最大值计算的维度。	取值范围在[-self.dim(), self.dim())。	-	-	-	-
  keepdim	输入	是否在输出张量中保留输入张量的维度。	-	-	-	-	-
  out	输出	输入Tensor。	-	INT32、INT64	ND	-	√
  workspaceSize	输出	返回需要在Device侧申请的workspace大小。	-	-	-	-	-
  executor	输出	返回op执行器，包含了算子计算流程。	-	-	-	-	-

  • Atlas 推理系列产品、Atlas 训练系列产品：不支持BFLOAT16数据类型。

• 返回值：

  aclnnStatus：返回状态码，具体参见aclnn返回码。

  第一段接口完成入参校验，出现以下场景时报错：

  返回码	错误码	描述
  ACLNN_ERR_PARAM_NULLPTR	161001	传入的self或out是空指针。
  ACLNN_ERR_PARAM_INVALID	161002	self、out的数据类型不在支持的范围之内。
  		self的shape超过8维。
  		dim不在合法范围内。

aclnnArgMax

• 参数说明：

  参数名	输入/输出	描述
  workspace	输入	在Device侧申请的workspace内存地址。
  workspaceSize	输入	在Device侧申请的workspace大小，由第一段接口aclnnArgMaxGetWorkspaceSize获取。
  executor	输入	op执行器，包含了算子计算流程。
  stream	输入	指定执行任务的Stream。

• 返回值：

  aclnnStatus：返回状态码，具体参见aclnn返回码。

约束说明

• 确定性计算：
  • aclnnArgMax默认确定性实现。

调用说明

调用方式	调用样例	说明
aclnn调用	test_aclnn_arg_max_v2	通过aclnnArgMax接口方式调用ArgMaxV2算子。