vaxpy

Description

Multiplies lhs by a scalar and adds rhs to the result element-wise.

Prototype

vaxpy(lhs, rhs, scalar)

Parameters

• lhs: a tvm.tensor for the left tensor
• rhs: a tvm.tensor for the right tensor
• scalar: a scalar for the coefficient by which lhs is multiplied
• The following data types are supported:

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

  Atlas Training Series Product: supports float16 and float32.

Note that:

• lhs and rhs must have the same data type and shape.
• If the data type of the scalar is different from that of the tensor, the data type of the scalar will be converted into the same one of the tensor for computation.

Returns

res_tensor: a tvm.tensor for the result tensor

Restrictions

None

Applicability

Atlas 200/300/500 Inference Product

Atlas Training Series Product

Example

from tbe import tvm
from tbe import dsl
shape = (1024,1024)
input_dtype = "float16"
data1 = tvm.placeholder(shape, name="data1", dtype=input_dtype)
data2 = tvm.placeholder(shape, name="data2", dtype=input_dtype)
scalar = tvm.const(2, dtype =input_dtype)
res = dsl.vaxpy(data1, data2, scalar)