TensorType

onnxscript/function_libs/torch_lib/ops/core.py

@@ -214,19 +214,21 @@
     by the scalar value and adds it to self.
     """
     # FIXME: Int to float
     self, tensor1, tensor2 = _type_promotion.promote_types(op, [self, tensor1, tensor2])
 
     return op.Add(self, op.Mul(op.Div(tensor1, tensor2), value))
 
 
 @torch_op("aten::addcmul", trace_only=True)
-def aten_addcmul(self: TReal, tensor1: TReal, tensor2: TReal, value: float = 1.0) -> TReal:
+def aten_addcmul(
+    self: TensorType, tensor1: TensorType, tensor2: TensorType, value: float = 1.0
+) -> TensorType:
     """addcmul(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor
 
     Performs the element-wise multiplication of tensor1 by tensor2, multiplies the
     result by the scalar value and adds it to self.
     """
     self, tensor1, tensor2 = _type_promotion.promote_types(op, [self, tensor1, tensor2])
 
     # Follow the order in https://github.com/pytorch/pytorch/blob/29e3fddb082b5a14262a7246bc62381a55199d45/aten/src/ATen/native/cpu/PointwiseOpsKernel.cpp#L47
     # TODO(#811): Understand fp16 accuracy issue
@@ -266,7 +268,7 @@
 
     Performs the outer-product of vectors vec1 and vec2 and adds it to the matrix input.
     """
     self, vec1, vec2 = _type_promotion.promote_types(op, [self, vec1, vec2])
     vec1_shape = op.Constant(value_ints=[-1, 1])
     vec2_shape = op.Constant(value_ints=[1, -1])
     vec1_reshaped = op.Reshape(vec1, vec1_shape)