更新时间:2021-03-18 GMT+08:00
分享

unsorted_segment_prod

功能说明

使用数组segment_ids对tensor进行分段求内积。假设输入为data,输出为output,则output[i] = product(data[j...]),其中“j...”是一个数组,“j...”中的元素j满足:segment_ids[j] == i。

product表示求内积,即data[j...]中所有元素相乘。

如果某个下标i在segment_ids中没有出现,则output[i] = init_value(默认值为0)。比如下图中,1在segment_ids中没有出现,则output[1] = 0。

如果segment_ids中某个值为负数,则对应位置的data中的值则丢弃。比如下图中,segment_ids[3] = -1,则data[3]的值被丢弃,不参与计算。

您可以在ATC安装目录下的“python/site-packages/te/te/lang/cce/te_compute/segment_compute.py”查看接口定义。

约束说明

昇腾310 AI处理器:支持的数据类型有float16、float32、int32

函数原型

te.lang.cce.unsorted_segment_prod(tensor, segment_ids, num_segments, init_value=0)

参数说明

  • tensor:输入tensor。
  • segment_ids:一维数组,对输入tensor进行分段的数组,其长度必须和输入tensor的第一维的长度相同。同时支持有序和无序。
  • num_segments:输出tensor的第一维的长度。其值必须大于等于segment_ids的最大值加1。
  • init_value:当segment_ids中某个下标不存在时,输出默认值0。

返回值

res_tensor:表示计算后的tensor。

调用示例

import tvm
import te.lang.cce 
shape = (5,1024) 
input_dtype = "float16"
data = tvm.placeholder(shape, name="data1", dtype=input_dtype) 
segment_ids = [1,1,4,5,5] 
num_segments = 6
res = te.lang.cce.unsorted_segment_prod(data, segment_ids, num_segments) 
# res.shape = (6,1024)
# res[0] = 0
# res[1] = (data[0] * data[1]) 
# res[2] = 0
# res[3] = 0
# res[4] = data[2]
# res[5] = (data[3] * data[4])
分享:

    相关文档

    相关产品

close