te.lang.cce.vaxpy(lhs, rhs, scalar)
将lhs中每个元素乘上标量scalar,再加上rhs中的对应元素,lhs和rhs tensor的数据类型要求一致,支持的类型:float16、float32。int8、uint8、int32类型会被转换为float16。
scalar的数据类型如果跟tensor不一致,会被转换成tensor的数据类型。
该接口在elewise_compute.py中定义。
参数说明
- lhs:左tensor,tvm.tensor类型。
- rhs:右tensor,tvm.tensor类型。
- scalar:lhs中元素要乘的系数,标量类型。
返回值
res_tensor:表示lhs * scalar + rhs,tvm.tensor类型。
调用示例
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)