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

    Atlas inference products: supports float16 and float32.

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

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

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

Applicable Models

Atlas training products

Atlas inference products

Atlas 200I/500 A2 inference products

Atlas A2 training products/Atlas A2 inference products

Atlas A3 training products/Atlas A3 inference products

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)