昇腾社区首页
中文
注册

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

输出