将输入张量沿指定维度切分成多个张量。
1 2 3 4 5 6 | struct SplitParam { int32_t splitDim = 0; int32_t splitNum = 2; SVector<int32_t> splitSizes = {}; uint8_t rsv[8] = {0}; }; |
成员名称 |
类型 |
默认值 |
描述 |
---|---|---|---|
splitDim |
int32_t |
0 |
指定切分的维度索引。 splitDim须位于输入张量x的维度范围内,即如果x的维度为xDim,则splitDim的取值范围为[-xDim, xDim - 1]。 当splitDim为负数时,其含义是从最高维度开始访问,如splitDim = -1,x维度数为dimNum,则拆分维度为dimNum - 1。 |
splitNum |
int32_t |
2 |
切分数量,当前支持2或3。 输入张量x的维度须能够被splitNum整除,且当splitNum = 3时输入x要求是float16或者bf16数据类型。 |
splitSizes |
SVector<int32_t> |
- |
指定每个输出tensor在切分维度上的大小,不传入此参数时使用等长切分,传入此参数时使用splitV不等长切分。 |
rsv[8] |
uint8_t |
{0} |
预留参数。 |
参数 |
维度 |
数据类型 |
格式 |
描述 |
---|---|---|---|---|
x |
[dim_0, ..., dim_splitDim,...,dim_n] |
float16/int64/bf16 |
ND |
输入,最高支持8维。 |
output1 |
|
float16/int64/bf16 |
ND |
输出,切分后的tensor。数据类型与x一致。 |
output2 |
|
float16/int64/bf16 |
ND |
输出,切分后的tensor。数据类型与x一致。 |
参数 |
维度 |
数据类型 |
格式 |
描述 |
---|---|---|---|---|
x |
[dim_0, ..., dim_splitDim,...,dim_n] |
float16/bf16 |
ND |
输入,最高支持8维。 |
output1 |
|
float16/bf16 |
ND |
输出,切分后的tensor。数据类型与x一致。 |
output2 |
|
float16/bf16 |
ND |
输出,切分后的tensor。数据类型与x一致。 |
output3 |
|
float16/bf16 |
ND |
输出,切分后的tensor。数据类型与x一致。 |
输入:
splitDim = 0 splitNum = 3 x = [3, 3, 3, 3, 3, 3, 3, 3, 3]
输出:
z = [3, 3, 3] z1 = [3, 3, 3] z2 = [3, 3, 3]