Help Center/ Atlas 500 Application/ TE API Reference/ Compute APIs/ te.lang.cce.broadcast(var, shape, output_dtype=None)
Updated on 2022-03-13 GMT+08:00

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

Broadcasts var to the tensor whose size is shape. The data type of the output is specified by output_dtype. var can be a scalar or tensor. The shape of var must have the same length as the second parameter shape . The size of each dimension must be either the same as that of shape, or be 1. When the size is 1, the dimension is broadcast to be the same as shape. For example, if the dimension of var is (2, 1, 64) and the value of shape is (2, 128, 64), the calculated dimension of var is (2, 128, 64). The supported types are float16, float32, and int32.

This API is defined in broadcast_compute.py.

Parameter Description

  • var: data to be broadcast, scalar or tensor type
  • shape: target shape for the broadcast operation
  • output_dtype: output data type. The default value is var.dtype.

Return Value

res_tensor: tensor obtained after var extension. The shape is specified by the parameter shape. The data type is output_dtype.

Calling Example

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)