SplitOperation
功能
对输入张量指定维度等分切成多个张量。
图1 SplitOperation


定义
struct SplitParam { int32_t splitDim = 0; int32_t splitNum = 2; };
成员
成员名称 |
描述 |
---|---|
splitDim |
指定切分的维度索引。 splitDim须位于输入张量x的维度范围内,即如果x的维度为xDim,则splitDim的取值范围为[-xDim, xDim - 1]。 当splitDim为负数时,其含义是从最高维度开始访问,如splitDim = -1,x维度数为dimNum,则拆分维度为dimNum - 1。 |
splitNum |
等分次数,当前支持2或3。 输入张量x的维度须能够被splitNum整除,且当splitNum = 3时输入x要求是float16或者bfloat16数据类型。 bfloat16类型数据仅支持Atlas 800I A2推理产品。 |
splitNum=2时输入输出
参数 |
维度 |
数据类型 |
格式 |
描述 |
---|---|---|---|---|
x |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/int64/bf16 |
ND |
输入,最高支持8维。 |
output1 |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/int64/bf16 |
ND |
输出 |
output2 |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/int64/bf16 |
ND |
输出 |
splitNum=3时输入输出
参数 |
维度 |
数据类型 |
格式 |
描述 |
---|---|---|---|---|
x |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/bf16 |
ND |
输入,最高支持8维。 |
output1 |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/bf16 |
ND |
输出 |
output2 |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/bf16 |
ND |
输出 |
output3 |
[-1,…,-1] -1表示当前维度的大小没有约束。 |
float16/bf16 |
ND |
输出 |