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 training product: supports float16 and float32.

    Atlas inference product: supports float16 and float32.

    Atlas 200I/500 A2 inference product: supports float16 and float32.

    Atlas A2 training product/Atlas A2 inference product: supports float16 and float32.

    Atlas A3 training product/Atlas A3 inference 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 training product

Atlas inference product

Atlas 200I/500 A2 inference product

Atlas A2 training product/Atlas A2 inference product

Atlas A3 training product/Atlas A3 inference 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)