broadcast

功能说明

把tensor var广播为形状为目标shape（第二个参数）的tensor，结果的数据类型由output_dtype指定。

如下A的shape为(2,1)即2行一列，广播为目标shape(2,3)即2行3列，则会将原来的一列扩展为一模一样的3列。

例如tensor var的shape为(2, 1, 64)，目标shape为(2, 128, 64)，调用broadcast接口的结果tensor的shape变为(2, 128, 64)。

函数原型

broadcast(var, shape, output_dtype=None)

参数说明

• var：需要broadcast的数据，标量或者tensor类型。

  昇腾310 AI处理器：支持的数据类型有float16、float32、int32。

  昇腾910 AI处理器：支持的数据类型有int8,uint8,float16,int16,uint16,float32,int32,uint32,int64和uint64。

  昇腾310P AI处理器：支持的数据类型有int8,uint8,float16,int16,uint16,float32,int32,uint32,int64和uint64。

• shape：目标shape，进行broadcast操作的目标shape。
• output_dtype：输出数据类型，默认值var.dtype。

返回值

res_tensor：由var扩展后得到的tensor，shape为参数指定的shape，数据类型为output_dtype。

约束说明

输入tensor var的shape长度需要小于或等于目标shape的长度，若var的shape长度小于目标shape，会在var前面自动补1，补齐为与目标shape长度相同。var每个维度的大小要么与目标shape相等，要么为1，为1的维度会被broadcast到与目标shape的对应维度一致。

支持的芯片型号

昇腾310 AI处理器

昇腾910 AI处理器

昇腾310P AI处理器

调用示例

如下代码所示，将shape为(1024,1)的tensor通过调用broadcast接口扩展至shape(1024,1024)。

from tbe import tvm
from tbe import dsl
outshape = (1024,1024)
shape = (1024,1)
input_dtype = "float16"
data = tvm.placeholder(shape, name="data", dtype=input_dtype)
res = dsl.broadcast(data, outshape)