feat(atenlib): add, sub, mul

ghstack-source-id: 89663f7
Pull Request resolved: #235

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -57,8 +57,9 @@ def aten_adaptive_max_pool1d(
 
 def aten_add(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:
     # add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
-
-    raise NotImplementedError()
+    if alpha != 1:
+        other = op.Mul(other, alpha)
+    return op.Add(self, other)
 
 
 def aten_addbmm(
@@ -3107,7 +3108,16 @@ def aten_msort(self: TensorType) -> TensorType:
 def aten_mul(self: TensorType, other: TensorType) -> TensorType:
     # mul.Tensor(Tensor self, Tensor other) -> Tensor
 
-    raise NotImplementedError()
+    return op.Mul(self, other)
+
+
+def aten_mul_bool(self: TensorType, other: TensorType) -> TensorType:
+    """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(
@@ -4337,7 +4347,10 @@ def aten_stft(
 def aten_sub(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:
     # sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
 
-    raise NotImplementedError()
+    if alpha != 1:
+        other = op.Mul(other, alpha)
+
+    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

@@ -7,6 +7,7 @@
 from typing import Callable, Collection, Iterable, Optional, Sequence
 
 import numpy as np
+import onnxruntime.capi.onnxruntime_pybind11_state
 import torch
 from torch.testing._internal import common_device_type, common_methods_invocations
 from torch.testing._internal.opinfo import core as opinfo_core
@@ -156,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),
@@ -179,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 #####################################################