Segment Compute APIs
Segment compute APIs are used to compute the sum, mean, inner product, maximum, minimum along the segments of a tensor.
unsorted_segment_sum
Description
Computes the sum along segments of a tensor by using the array segment_ids. Assuming that the input is data and the output is output, then output[i] = sum(data[j...]), where j... is an array, and element j meets the following requirement: segment_ids[j] == i.
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 API is defined in python/site-packages/te/lang/cce/te_compute/segment_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16, float32, int32
Prototype
te.lang.cce.unsorted_segment_sum(tensor, segment_ids, num_segments, init_value=0)
Parameters
- tensor: input tensor
- segment_ids: a 1D array used to segment the input tensor. Has the same length as that 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 segment_ids plus 1.
- init_value: default element value for segment_ids. It is determined based on the implementation of the operator. Defaults to 0.
Returns
res_tensor: tensor after computation
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_sum(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]
unsorted_segment_mean
Description
Computes the mean along segments of a tensor by using the array segment_ids. Assuming that the input is data and the output is output, then output[i] = (1/len(j...)) sum(data[j...]), where j... is an array, and element j meets the following requirement: segment_ids[j] == i.
If a subscript i does not appear in segment_ids, then output[i] = init_value (the default value is 0). 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 API is defined in python/site-packages/te/lang/cce/te_compute/segment_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16, float32, int32
Prototype
te.lang.cce.unsorted_segment_mean(tensor, segment_ids, num_segments, init_value=0)
Parameters
- tensor: input tensor
- segment_ids: a 1D array used to segment the input tensor. Has the same length as that 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 segment_ids plus 1.
- init_value: default element value for segment_ids. It is determined based on the implementation of the operator. Defaults to 0.
Returns
res_tensor: tensor after computation
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,5,5,5] num_segments = 6 res = te.lang.cce.unsorted_segment_mean(data, segment_ids, num_segments) # res.shape = (6,1024) # res[0] = 0 # res[1] = (data[0] + data[1]) / 2 # res[2] = 0 # res[3] = 0 # res[4] = 0 # res[5] = (data[2] + data[3] + data[4]) / 3
unsorted_segment_prod
Description
Computes the inner product along segments of a tensor by using the array segment_ids. 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 (the default value is 0). 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 API is defined in python/site-packages/te/lang/cce/te_compute/segment_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16, float32, int32
Prototype
te.lang.cce.unsorted_segment_prod(tensor, segment_ids, num_segments, init_value=0)
Parameters
- tensor: input tensor
- segment_ids: a 1D array used to segment the input tensor. Has the same length as that 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 segment_ids plus 1.
- init_value: default element value for segment_ids, fixed at 0.
Returns
res_tensor: tensor after computation
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] = 0 # res[1] = (data[0] * data[1]) # res[2] = 0 # res[3] = 0 # res[4] = data[2] # res[5] = (data[3] * data[4])
unsorted_segment_min
Description
Computes the minimum along segments of a tensor by using the array segment_ids. Assuming that the input is data and the output is output, then output[i] = min(data[j...]), where j... is an array, and element j meets the following requirement: segment_ids[j] == i.
If a subscript i does not appear in segment_ids, then output[i] = init_value (the default value is 0). 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 API is defined in python/site-packages/te/lang/cce/te_compute/segment_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16, float32, int32
Prototype
te.lang.cce.unsorted_segment_min(tensor, segment_ids, num_segments, init_value=0)
Parameters
- tensor: input tensor
- segment_ids: a 1D array used to segment the input tensor. Has the same length as that 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 segment_ids plus 1.
- init_value: default element value for segment_ids, fixed at 0.
Returns
res_tensor: tensor after computation
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_min(data, segment_ids, num_segments) # res.shape = (6,1024) # res[0] = 0 # res[1] = min(data[0], data[1]) # res[2] = 0 # res[3] = 0 # res[4] = data[2] # res[5] = min(data[3], data[4])
unsorted_segment_max
Description
Computes the maximum along segments of a tensor by using the array segment_ids. Assuming that the input is data and the output is output, then output[i] = max(data[j...]), where j... is an array, and element j meets the following requirement: segment_ids[j] == i.
If a subscript i does not appear in segment_ids, then output[i] = init_value (the default value is 0). 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 API is defined in python/site-packages/te/lang/cce/te_compute/segment_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16, float32, int32
Prototype
te.lang.cce.unsorted_segment_max(tensor, segment_ids, num_segments, init_value=0)
Parameters
- tensor: input tensor
- segment_ids: a 1D array used to segment the input tensor. Has the same length as that 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 segment_ids plus 1.
- init_value: default element value for segment_ids. It is determined based on the implementation of the operator. Defaults to 0.
Returns
res_tensor: tensor after computation
Example
import tvm import te.lang.cce shape = (5,1024) input_dtype = "float16" data = tvm.placeholder(shape, name="data1",type=input_dtype) segment_ids = [1,1,4,5,5] num_segments = 6 res = te.lang.cce.unsorted_segment_max(data,segment_ids,num_segments) # res.shape = (6,1024) # res[0] = 0 # res[1] = max(data[0], data[1]) # res[2] = 0 # res[3] = 0 # res[4] = data[2] # res[5] = max(data[3], data[4])
