Help Center/ Atlas 300 Application (Model 3000)/ TE API Reference/ Compute APIs/ te.lang.cce.unsorted_segment_prod(tensor, segment_ids, num_segments, init_value=0)
Updated on 2022-03-13 GMT+08:00
View PDF
Share

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

Uses the array segment_ids to calculate the inner product of tensors by segment. Assuming that the input is data and the output is output, then output[i] = product(data[j...]), where j... is an array, and element j meets the following requirement: segment_ids[j] == i.

The parameter product indicates the inner product, that is, all elements in data[j...] are multiplied by each other.

If a subscript i does not appear in segment_ids, then output[i] = init_value. For example, in the following figure, if 1 does not appear in segment_ids, then output[1] = 0.

If a value in segment_ids is a negative number, the value of data in the corresponding position is discarded. For example, in the following figure, segment_ids[3] = –1, indicating that the value of data[3] is discarded and is not involved in the calculation.

The length of segment_ids must be the same as the length of the first dimension of data. The value of num_segments must be greater than or equal to the value of segment_ids plus 1.

The supported types are float16, float32, and int32.

This API is defined in segment_compute.py.

Parameter Description

  • tensor: input tensor, which is of the float16, float32, or int32 type
  • segment_ids: one-dimensional array. This array is used to segment the input tensor. The length of the array must be the same as the length of the first dimension of the input tensor. The array can be sequenced or unsequenced.
  • num_segments: length of the first dimension of the output tensor. Its value must be greater than or equal to the value of segment_ids plus 1.
  • init_value: default value of the output when a subscript in segment_ids does not exist. It is determined according to the implementation of the operator. The default value is 0.

Return Value

res_tensor: tensor after calculation

Calling Example

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] = 1
# res[1] = (data[0] * data[1]) 
# res[2] = 1
# res[3] = 1
# res[4] = data[2]
# res[5] = (data[3] * data[4])
Parent topic: Compute APIs