SplitOperation
Description
Equally divides a specific dimension of the input tensor to multiple tensors.
Operator Context
Operator Function Implementation
- Parameters:
Parameter
Description
x
Input tensor, that is, the tensor to be split.
splitNum
The value is an integer, indicating the number of splits. Currently, the value can be 2 or 3.
splitDim
Optional. The default value is 0, indicating that splitting is performed along this dimension. By default, splitting is performed along the first dimension (that is, the batch dimension).
- Returns:
Returns several tensors. These tensors are split from the original tensor in a specified manner.
Definition
struct SplitParam {
int32_t splitDim = 0;
int32_t splitNum = 2;
SVector<int32_t> splitSizes = {};
uint8_t rsv[8] = {0};
};
Parameters
Member |
Type |
Default Value |
Description |
|---|---|---|---|
splitDim |
int32_t |
0 |
Index of the dimension to be split. splitDim must be within the dimension range of input x. For example, if the dimension of x is xDim, the value range of splitDim is [–xDim, xDim – 1]. If splitDim is a negative number, the access starts from the highest dimension. For example, if splitDim is –1 and the number of dimensions of x is dimNum, the number of dimensions to be split is dimNum – 1. |
splitNum |
int32_t |
2 |
Number of splits. Currently, the value can be 2 or 3. The dimension of the input tensor x must be exactly divided by splitNum. When splitNum is 3, x must be of the float16 or bf16 data type. |
splitSizes |
SVector<int32_t> |
- |
Size of each output tensor in the split dimension. If this parameter is not passed, equal-length split is used. If this parameter is passed, splitV is used. |
rsv[8] |
uint8_t |
{0} |
Reserved |
Input and Output (splitNum=2)
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
x |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/int64/bf16 |
ND |
Input, up to 8 dimensions. |
output1 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/int64/bf16 |
ND |
Output, which has the same dimension as x. |
output2 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/int64/bf16 |
ND |
Output, which has the same dimension as x. |
Input and Output (splitNum=3)
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
x |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/bf16 |
ND |
Input, up to 8 dimensions. |
output1 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/bf16 |
ND |
Output, which has the same dimension as x. |
output2 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/bf16 |
ND |
Output, which has the same dimension as x. |
output3 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16/bf16 |
ND |
Output, which has the same dimension as x. |
Restrictions
- splitDim must be within the dimension range of input x. For example, if the dimension of x is xDim, the value range of splitDim is [–xDim, xDim – 1]. If splitDim is a negative number, the access starts from the highest dimension. For example, if splitDim is –1 and the number of dimensions of x is dimNum, the number of dimensions to be split is dimNum – 1. Negative indexes are not supported when unequal-length splitting is used.
- The dimension of the input tensor x must be divisible by splitNum.
- When splitNum is 3, the data type of x must be float16 or bfloat16.
- Each element in splitSizes must be greater than or equal to 1. The sum of elements in splitSizes is equal to the size of the split dimension.
API Calling Example
Input:
splitDim = 0 splitNum = 3 x = [3, 3, 3, 3, 3, 3, 3, 3, 3]
Output:
z = [3, 3, 3] z1 = [3, 3, 3] z2 = [3, 3, 3]