Updated on 2022-03-13 GMT+08:00

te.lang.cce.vaxpy(lhs, rhs, scalar)

Multiplies each element in lhs by a scalar and adds the corresponding element in rhs. The data types of lhs and rhs must be the same. The supported types are float16 and float32. The data types int8, uint8, and int32 are converted to float16.

If the data type of the scalar is different from that of the tensor, the data type of the tensor prefers.

This API is defined in elewise_compute.py.

Parameter Description

  • lhs: left tensor, tvm.tensor type
  • rhs: right tensor, tvm.tensor type
  • scalar: coefficient by which the element in lhs is multiplied, scalar type

Return Value

res_tensor: lhs x scalar + rhs, tvm.tensor type

Calling Example

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 = te.lang.cce.vaxpy(data1, data2, scalar)