# 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.

**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.

**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.

** 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.

**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.

**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])