SplitOperation

Hardware Model	Supported
Atlas 350 accelerator cards	√
Atlas A3 inference products/Atlas A3 training products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas training products	√
Atlas inference products	√
Atlas 200I/500 A2 inference products	√

Splits an input tensor into multiple tensors along a specified dimension.

Figure 1 SplitOperation

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).
splitSizes	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, unequal-length split is used.

Returns:
Returns several tensors. These tensors are split from the original tensor in a specified manner.

struct SplitParam {
    int32_t splitDim = 0;
    int32_t splitNum = 2;
    SVector<int32_t> splitSizes = {};
    uint8_t rsv[8] = {0};
};

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, unequal-length split is used.
rsv[8]	uint8_t	{0}	Reserved

Parameter	Dimension	Data Type	Format	Description
x	[dim_0, ..., dim_splitDim, ..., dim_n]	float16/int64/bf16	ND	Input, up to 8 dimensions.
output1	[dim_0, ..., dim_splitDim_1, ..., dim_n] If splitSizes is empty: dim_splitDim_1 = dim_splitDim/2 If splitSizes is not empty: dim_splitDim_2 = splitSizes[0]	float16/int64/bf16	ND	Output, tensor after splitting. The data type is the same as that of x.
output2	[dim_0, ..., dim_splitDim_2, ..., dim_n] If splitSizes is empty: dim_splitDim_2 = dim_splitDim/2. If splitSizes is not empty: dim_splitDim_2 = splitSizes[1].	float16/int64/bf16	ND	Output, tensor after splitting. The data type is the same as that of x.

Parameter	Dimension	Data Type	Format	Description
x	[dim_0, ..., dim_splitDim, ..., dim_n]	float16/bf16	ND	Input, up to 8 dimensions.
output1	[dim_0, ..., dim_splitDim_1, ..., dim_n] If splitSizes is empty: dim_splitDim_1 = dim_splitDim/2 If splitSizes is not empty: dim_splitDim_1 = splitSizes[0]	float16/bf16	ND	Output, tensor after splitting. The data type is the same as that of x.
output2	[dim_0, ..., dim_splitDim_2, ..., dim_n] If splitSizes is empty: dim_splitDim_2 = dim_splitDim/3 If splitSizes is not empty: dim_splitDim_2 = splitSizes[1].	float16/bf16	ND	Output, tensor after splitting. The data type is the same as that of x.
output3	[dim_0, ..., dim_splitDim_3, ..., dim_n] If splitSizes is empty: dim_splitDim_3 = dim_splitDim/3 If splitSizes is not empty: dim_splitDim_3 = splitSizes[2]	float16/bf16	ND	Output, tensor after splitting. The data type is the same as that of x.

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.
If splitSizes is not empty, the dimension is splitNum. Each element must be greater than or equal to 1. The sum of elements is equal to the size of dim_splitDim.
If splitSizes is empty, the dimension of the input tensor x must be divisible by splitNum.
If splitNum is 3, the data type of x must be float16 or bf16.
The bf16 data type is supported only by Atlas 350 accelerator cards, Atlas A2 training products/Atlas A2 inference products, and Atlas A3 inference products/Atlas A3 training products.

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]

Parent topic: atb/infer_op_params.h