文档首页 > > API参考> TE API参考> compute接口>

te.lang.cce.broadcast(var, shape, output_dtype=None)

te.lang.cce.broadcast(var, shape, output_dtype=None)

分享
更新时间:2020/05/11 GMT+08:00

把var broadcast为大小为shape的tensor,结果的数据类型由output_dtype指定,var可以是标量,或者是一个tensor,要求var的shape与第二个参数shape的长度一致,每个维度的大小要么与shape相等,要么为1,为1的维度会被broadcast到与shape一致。例如var的维度为(2,1,64),shape为(2,128,64),运算结果var的维度变为(2,128,64)。支持的类型:float16、float32、int32。

该接口在broadcast_compute.py中定义。

参数说明

  • var:需要broadcast的数据,标量或者tensor类型。
  • shape:目标shape,进行broadcast操作的目标shape。
  • output_dtype:输出数据类型,默认值var.dtype。

返回值

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

调用示例

outshape = (1024,1024)
shape = (1024,1)
input_dtype = "float16"
data = tvm.placeholder(shape, name="data", dtype=input_dtype)
res = te.lang.cce.broadcast(data, outshape)
分享:

    相关文档

    相关产品

文档是否有解决您的问题?

提交成功!非常感谢您的反馈,我们会继续努力做到更好!
反馈提交失败,请稍后再试!

*必选

请至少选择或填写一项反馈信息

字符长度不能超过200

提交反馈 取消

如您有其它疑问,您也可以通过华为云社区论坛频道来与我们联系探讨

智能客服提问云社区提问