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
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)
Parent topic: Math Compute APIs