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
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