feat(atenlib): ops 3/n

A bunch of matmul ops

ghstack-source-id: da6585a
Pull Request resolved: #255

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -22,28 +22,22 @@
 from onnxscript.onnx_types import TensorType
 
 
-def aten_abs(self: TensorType) -> TensorType:
+def aten_abs(self):
     # abs(Tensor self) -> Tensor
 
-    raise NotImplementedError()
-
-
-def aten_absolute(self: TensorType) -> TensorType:
-    # absolute(Tensor self) -> Tensor
-
-    raise NotImplementedError()
+    return op.Abs(self)
 
 
-def aten_acos(self: TensorType) -> TensorType:
+def aten_acos(self):
     # acos(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Acos(self)
 
 
-def aten_acosh(self: TensorType) -> TensorType:
+def aten_acosh(self):
     # acosh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Acosh(self)
 
 
 def aten_adaptive_avg_pool1d(self: TensorType, output_size: Sequence[int]) -> TensorType:
@@ -91,12 +85,13 @@ def aten_addcmul(
     raise NotImplementedError()
 
 
-def aten_addmm(
-    self: TensorType, mat1: TensorType, mat2: TensorType, beta: float = 1, alpha: float = 1
-) -> TensorType:
+def aten_addmm(self, mat1, mat2, beta: float = 1, alpha: float = 1):
     # addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor
 
-    raise NotImplementedError()
+    mat1_mat2 = op.MatMul(mat1, mat2)
+    scaled_mat1_mat2 = op.Mul(mat1_mat2, alpha)  # type: ignore[arg-type]
+    scaled_self = op.Mul(self, beta)  # type: ignore[arg-type]
+    return op.Add(scaled_self, scaled_mat1_mat2)  # type: ignore[arg-type]
 
 
 def aten_addmv(
@@ -611,10 +606,10 @@ def aten_block_diag(tensors: Sequence[TensorType]) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_bmm(self: TensorType, mat2: TensorType) -> TensorType:
+def aten_bmm(self, mat2):
     # bmm(Tensor self, Tensor mat2) -> Tensor
 
-    raise NotImplementedError()
+    return op.MatMul(self, mat2)
 
 
 def aten_broadcast_tensors(tensors: Sequence[TensorType]) -> TensorType:

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -161,7 +161,10 @@ def wrapped(fn):
 # Ops to be tested for numerical consistency between onnx and pytorch
 # Find the names of the OpInfos in torch/testing/_internal/common_methods_invocations.py
 OPINFO_FUNCTION_MAPPING: dict[str, Callable[..., Any]] = {
+    "abs": core_ops.aten_abs,
     "add": core_ops.aten_add,
+    "addmm": core_ops.aten_addmm,
+    "bmm": core_ops.aten_bmm,
     "clamp_max": core_ops.aten_clamp_max_tensor,
     "clamp_min": core_ops.aten_clamp_min_tensor,
     "clamp": core_ops.aten_clamp,
@@ -187,6 +190,22 @@ def wrapped(fn):
 
 EXPECTED_SKIPS_OR_FAILS = (
     xfail("add", dtypes=BOOL_TYPES, reason="Add is not defined on bool tensors"),
+    xfail(
+        "addmm",
+        dtypes=[torch.uint8, torch.int8, torch.int16],
+        reason="MatMul is not defined on int16/int8/uint8 tensors",
+    ),
+    xfail(
+        "addmm",
+        variant_name="decomposed",
+        dtypes=[torch.uint8, torch.int8, torch.int16],
+        reason="MatMul is not defined on int16/int8/uint8 tensors",
+    ),
+    xfail(
+        "bmm",
+        dtypes=[torch.uint8, torch.int8, torch.int16],
+        reason="MatMul is not defined on int16/int8/uint8 tensors",
+    ),
     skip("clamp", reason="Enable when onnxscript errors are fixed"),
     xfail("clamp_max", dtypes=BOOL_TYPES, reason="Min is not defined on bool tensors"),
     xfail("clamp_min", dtypes=BOOL_TYPES, reason="Max is not defined on bool tensors"),