feat(atenlib): add, sub, mul

onnxscript/function_libs/torch_aten/ops/core.py

@@ -17,7 +17,7 @@
 
 from typing import Any, Optional, Sequence
 
-from onnxscript import INT64
+from onnxscript import BOOL, INT64
 from onnxscript.onnx_opset import default_opset as op
 from onnxscript.onnx_types import TensorType
 
@@ -60,10 +60,11 @@ def aten_adaptive_max_pool1d(
     raise NotImplementedError()
 
 
-def aten_add(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:
+def aten_add(self, other, alpha: float = 1) -> TensorType:
     # add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
-
-    raise NotImplementedError()
+    if alpha != 1:
+        other = op.Mul(other, alpha)  # type: ignore[arg-type]
+    return op.Add(self, other)
 
 
 def aten_addbmm(
@@ -3109,10 +3110,19 @@ def aten_msort(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_mul(self: TensorType, other: TensorType) -> TensorType:
+def aten_mul(self, other) -> TensorType:
     # mul.Tensor(Tensor self, Tensor other) -> Tensor
 
-    raise NotImplementedError()
+    return op.Mul(self, other)
+
+
+def aten_mul_bool(self: BOOL, other: BOOL) -> BOOL:
+    """ONNX Mul doesn't support Boolean, so use And as an equivalent operator."""
+
+    # TODO(justinchuby): Handle cases where type reconcilation is not enough,
+    # since different ONNX operators are used based on different data types.
+
+    return op.And(self, other)
 
 
 def aten_multinomial(
@@ -4339,10 +4349,13 @@ def aten_stft(
     raise NotImplementedError()
 
 
-def aten_sub(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:
+def aten_sub(self, other, alpha: float = 1) -> TensorType:
     # sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
 
-    raise NotImplementedError()
+    if alpha != 1:
+        other = op.Mul(other, alpha)  # type: ignore[arg-type]
+
+    return op.Sub(self, other)
 
 
 def aten_subtract(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -157,14 +157,19 @@ def wrapped(fn):
 
 # Ops to be tested for numerical consistency between onnx and pytorch
 OPINFO_FUNCTION_MAPPING = {
+    "add": core_ops.aten_add,
+    "mul": core_ops.aten_mul,
     "nn.functional.elu": nn_ops.aten_elu,
     "nn.functional.relu6": nn_ops.aten_relu6,
     "nn.functional.selu": core_ops.aten_selu,
+    "sub": core_ops.aten_sub,
 }
 
 TESTED_OPS = frozenset(OPINFO_FUNCTION_MAPPING)
 
 EXPECTED_SKIPS_OR_FAILS = (
+    xfail("add", dtypes=BOOL_TYPES, reason="Add is not defined on bool tensors"),
+    xfail("mul", dtypes=BOOL_TYPES, reason="Mul is not defined on bool tensors"),
     xfail(
         "nn.functional.elu",
         dtypes=dtypes_except(torch.float16, torch.float32),
@@ -180,6 +185,7 @@ def wrapped(fn):
         dtypes=dtypes_except(torch.float16, torch.float32),
         reason="ONNX Runtime doesn't support float64 for Selu",
     ),
+    xfail("sub", dtypes=BOOL_TYPES, reason="Sub is not defined on bool tensors"),
 )
 # END OF SECTION TO MODIFY #####################################################