# pypto.topk【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√功能说明获取最后一个维度的前k个最大值或最小值及其对应的索引。如果输入是向量则在向量中找到前k个最大值或最小值及其对应的索引如果输入是矩阵则沿最后一个维度计算每行中前k个最大值或最小值及其对应的索引。如下图所示对Shape为(4, 32)的二维矩阵进行排序k设置为1输出结果为[[32] [32] [32] [32]]。函数原型topk(input: Tensor, k: int, dim: Optional[int] None, largest: bool True, algo: TopKAlgo TopKAlgo.MERGE_SORT) - Tuple[Tensor, Tensor]参数说明参数名输入/输出说明input输入源操作数。支持的类型为Tensor。Tensor支持的数据类型为- MERGE_SORT: DT_FP32。- RADIX_SELECT: DT_BF16DT_FP16DT_FP32。不支持空TensorShape仅支持1-4维Shape Size不大于2147483647即INT32_MAX。k输入返回元素的数量。k的大小应该满足1 k input.shape[dim]。dim输入指定排序的维度。目前仅支持按最后一个维度排序即dim -1或dim input.shape.size() - 1。largest输入如果为True返回最大元素。如果为False返回最小元素。algo输入算法枚举类型用以控制TopK计算的流程具体定义为TopKAlgo 。默认为 MERGE_SORT归并排序算法。返回值说明返回一个命名元组(values, indices)其中包含input在指定维度dim下每行中最大或最小的k个元素的值和索引。约束说明只支持对尾轴进行topk操作TileShape尾轴32bytes对齐(TileShape[-1]*4 % 32 0)选用MERGE_SORT算法时TileShape尾轴需要小于22KB(TileShape[-1]*4 22KB)选用RADIX_SELECT算法时记TileShape尾轴为tile则需要临时空间2*tile*sizeof(srcType)6*tile1024max(1024, 8*tile)临时空间加上输入输出的tile块不能超过UB大小k TileShape[-1] k input.shape[-1]RADIX_SELECT算法仅支持Ascend 950PR/Ascend 950DT调用示例TileShape设置示例说明调用该operation接口前应通过set_vec_tile_shapes设置TileShape。TileShape维度应和输入input一致。示例1输入input shape为[m, n, p]dim为2largest为True输出为[m, n, k], TileShape设置为[m1, n1, p1], 则m1, n1, p1分别用于切分m, n, p轴。p1必须大于等于kk轴不支持切分必须保证全载。pypto.set_vec_tile_shapes(4, 16, 32)接口调用示例x pypto.tensor([2, 3], pypto.DT_FP32) y pypto.topk(x, 2, -1, True, pypto.TopKAlgo.MERGE_SORT)结果示例如下输入数据x: [[1.0 2.0 3.0], [1.0 2.0 3.0]] 输出数据y[0]: [[3.0 2.0], [3.0 2.0]] 输出数据y[1]: [[2, 1], [2, 1]]【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考