CANN/ops-nn Index算子

📅 2026/7/6 16:18:43
CANN/ops-nn Index算子
Index【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品×Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×功能说明算子功能根据indices从输入张量x中取出指定位置的数据生成输出张量y。接口形态当前实验目录提供aclnnIndexGetWorkspaceSize/aclnnIndex接口内部将输入和索引转为连续张量后执行 Index 计算。执行路径普通INT32/INT64索引走IndexAiCoreBOOL 索引、索引数量超过输入 rank 等特殊路径走IndexAiCpu或直接返回。参数说明参数名输入/输出描述数据类型数据格式维度x输入待索引的输入张量对应 aclnn 接口中的 self。FLOAT、DOUBLE、FLOAT16、BFLOAT16、INT16、INT32、INT64、INT8、UINT8、BOOLND0-8indexed_sizes输入内部元数据记录被索引维度信息由 aclnn 接口根据 self 和 indices 生成。INT64ND1indexed_strides输入内部元数据记录输入张量 stride 或 BOOL 索引输出形状信息由 aclnn 接口生成。INT64ND1indices动态输入索引张量列表。普通 AICore 路径支持 INT32、INT64aclnn 参数检查接受 BOOLBOOL 索引走 AICPU 路径。INT32、INT64、BOOLND0-8y输出根据索引取出后的输出张量数据类型需与 x 一致。FLOAT、DOUBLE、FLOAT16、BFLOAT16、INT16、INT32、INT64、INT8、UINT8、BOOLND0-8约束说明仅支持 ND 格式不支持私有格式输入、索引或输出。x、indices中每个索引张量、y的 rank 均不超过 8。普通索引路径要求indices数量在[1, x rank]范围内且各索引张量 shape 可广播输出 shape 需与接口推导结果一致。空 Tensor 和标量输入的单索引场景由 aclnn 接口做特殊处理。AICore 二进制当前配置覆盖ascend910b和ascend950。调用说明调用方式接口/文件说明aclnn 调用op_api/aclnn_index_v2.haclnnIndexGetWorkspaceSize/aclnnIndex当前目录未提供examples/test_aclnn_index.cpp可通过 aclnn 接口头文件和 ATK 用例进行验证。Level0 内部调用op_api/experimental_index_l0.hIndexAiCore/IndexAiCpuaclnn 接口根据索引类型和形状选择 AICore 或 AICPU 路径。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考