broadcast

Description

Broadcasts var to the target shape (the second parameter). The data type of the result is specified by output_dtype.

For example, tensor A (2, 1) has two rows and one column. It can be broadcast to the target shape (2, 3) — two rows and three columns. In this case, the original column is broadcast to three identical columns.

For example, if the shape of var is (2, 1, 64) and the target shape is (2, 128, 64), the shape of the result tensor after the broadcast API call is (2, 128, 64).

Prototype

broadcast(var, shape, output_dtype=None)

Parameters

• var: a scalar or tensor for data to be broadcast.

  Atlas 200/300/500 Inference Product: supports float16, float32, and int32.

  Atlas Training Series Product: supports int8, uint8, float16, int16, uint16, float32, int32, uint32, int64, and uint64.

• shape: target shape for the broadcast operation.
• output_dtype: output data type. Defaults to var.dtype.

Returns

res_tensor: a result tensor broadcast from var, which has the target shape and is of type output_dtype.

Restrictions

The shape length of var of the input tensor must be less than or equal to that of the target shape. Otherwise, leading dimensions sized 1 will be added to var to make the target shape length. The size of each dimension of var must be consistent with that of the target shape or equal to 1. Any dimension sized 1 will be broadcast to the corresponding dimension of the target shape.

Applicability

Atlas 200/300/500 Inference Product

Atlas Training Series Product

Example

As shown in the following code, the shape (1024, 1) of a tensor is broadcast to shape (1024, 1024) by calling the broadcast API.

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)