feat(atenlib): ops 4/n

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()
-
+    return op.Abs(self)
 
-def aten_absolute(self: TensorType) -> TensorType:
-    # absolute(Tensor self) -> Tensor
-
-    raise NotImplementedError()
 
-
-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(
@@ -337,22 +332,22 @@ def aten_as_strided_scatter(
     raise NotImplementedError()
 
 
-def aten_asin(self: TensorType) -> TensorType:
+def aten_asin(self):
     # asin(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Asin(self)
 
 
-def aten_asinh(self: TensorType) -> TensorType:
+def aten_asinh(self):
     # asinh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Asinh(self)
 
 
-def aten_atan(self: TensorType) -> TensorType:
+def aten_atan(self):
     # atan(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Atan(self)
 
 
 def aten_atan2(self: TensorType, other: TensorType) -> TensorType:
@@ -361,10 +356,10 @@ def aten_atan2(self: TensorType, other: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_atanh(self: TensorType) -> TensorType:
+def aten_atanh(self):
     # atanh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Atanh(self)
 
 
 def aten_atleast_1d(self: TensorType) -> TensorType:
@@ -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:
@@ -675,16 +670,10 @@ def aten_cdist(
     raise NotImplementedError()
 
 
-def aten_ceil(self: TensorType) -> TensorType:
+def aten_ceil(self):
     # ceil(Tensor self) -> Tensor
 
-    raise NotImplementedError()
-
-
-def aten_celu(self: TensorType, alpha: float = 1.0) -> TensorType:
-    # celu(Tensor self, Scalar alpha=1.0) -> Tensor
-
-    raise NotImplementedError()
+    return op.Ceil(self)
 
 
 def aten_chain_matmul(matrices: Sequence[TensorType]) -> TensorType:
@@ -790,14 +779,6 @@ def aten_clamp_min_tensor(self, min_):
     return op.Max(self, min_)
 
 
-def aten_clip(
-    self: TensorType, min: Optional[float] = None, max: Optional[float] = None
-) -> TensorType:
-    # clip(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor
-
-    raise NotImplementedError()
-
-
 def aten_clone(self: TensorType, memory_format: Optional[str] = None) -> TensorType:
     # clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor
 
@@ -1036,16 +1017,16 @@ def aten_corrcoef(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_cos(self: TensorType) -> TensorType:
+def aten_cos(self):
     # cos(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Cos(self)
 
 
-def aten_cosh(self: TensorType) -> TensorType:
+def aten_cosh(self):
     # cosh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Cosh(self)
 
 
 def aten_cosine_embedding_loss(
@@ -1411,10 +1392,10 @@ def aten_divide(self: TensorType, other: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_dot(self: TensorType, tensor: TensorType) -> TensorType:
+def aten_dot(self, tensor):
     # dot(Tensor self, Tensor tensor) -> Tensor
 
-    raise NotImplementedError()
+    return op.MatMul(self, tensor)
 
 
 def aten_dropout(input: TensorType, p: float, train: bool) -> TensorType:
@@ -1551,16 +1532,18 @@ def aten_erfinv(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_exp(self: TensorType) -> TensorType:
+def aten_exp(self):
     # exp(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Exp(self)
 
 
-def aten_exp2(self: TensorType) -> TensorType:
+def aten_exp2(self):
     # exp2(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    two = op.Constant(value_int=2)
+    two = op.CastLike(two, self)  # type: ignore[arg-type]
+    return op.Pow(two, self)  # type: ignore[arg-type]
 
 
 def aten_expand(self: TensorType, size: INT64, implicit: bool = False) -> TensorType:
@@ -2680,10 +2663,10 @@ def aten_masked_select_backward(
     raise NotImplementedError()
 
 
-def aten_matmul(self: TensorType, other: TensorType) -> TensorType:
+def aten_matmul(self, other):
     # matmul(Tensor self, Tensor other) -> Tensor
 
-    raise NotImplementedError()
+    return op.MatMul(self, other)
 
 
 def aten_matmul_backward(
@@ -3080,10 +3063,11 @@ def aten_mkldnn_max_pool3d_backward(
     raise NotImplementedError()
 
 
-def aten_mm(self: TensorType, mat2: TensorType) -> TensorType:
+def aten_mm(self, mat2):
     # mm(Tensor self, Tensor mat2) -> Tensor
 
-    raise NotImplementedError()
+    # TODO(justinchuby): Specify type conversion for uint8/int8/int16
+    return op.MatMul(self, mat2)
 
 
 def aten_mode(
@@ -3463,16 +3447,13 @@ def aten_nuclear_norm(self: TensorType, keepdim: bool = False) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_numpy_T(self: TensorType) -> TensorType:
-    # numpy_T(Tensor(a) self) -> Tensor(a)
-
-    raise NotImplementedError()
-
-
-def aten_ones(size: INT64) -> TensorType:
+def aten_ones(size: INT64, dtype: int = -1) -> TensorType:
     # ones(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
 
-    raise NotImplementedError()
+    one = op.Constant(value_float=1)
+    if dtype != -1:
+        one = op.Cast(one, to=dtype)  # type: ignore[arg-type]
+    return op.Expand(one, size)  # type: ignore[arg-type]
 
 
 def aten_ones_like(self, dtype: int = -1):
@@ -4212,22 +4193,16 @@ def aten_signbit(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_sin(self: TensorType) -> TensorType:
+def aten_sin(self):
     # sin(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Sin(self)
 
 
-def aten_sinc(self: TensorType) -> TensorType:
-    # sinc(Tensor self) -> Tensor
-
-    raise NotImplementedError()
-
-
-def aten_sinh(self: TensorType) -> TensorType:
+def aten_sinh(self):
     # sinh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Sinh(self)
 
 
 def aten_slice(
@@ -4461,7 +4436,13 @@ def aten_symeig(
 def aten_t(self: TensorType) -> TensorType:
     # t(Tensor(a) self) -> Tensor(a)
 
-    raise NotImplementedError()
+    # TODO(justinchuby): Make rank a function
+    rank = op.Size(op.Shape(self))  # type: ignore[arg-type]
+    if rank == 0 or rank == 1:  # pylint: disable=consider-using-in
+        result = self
+    else:
+        result = op.Transpose(self, perm=[1, 0])  # type: ignore[arg-type]
+    return result
 
 
 def aten_t_copy(self: TensorType) -> TensorType:
@@ -4484,16 +4465,16 @@ def aten_take_along_dim(
     raise NotImplementedError()
 
 
-def aten_tan(self: TensorType) -> TensorType:
+def aten_tan(self):
     # tan(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Tan(self)
 
 
-def aten_tanh(self: TensorType) -> TensorType:
+def aten_tanh(self):
     # tanh(Tensor self) -> Tensor
 
-    raise NotImplementedError()
+    return op.Tanh(self)
 
 
 def aten_tensordot(
@@ -4606,6 +4587,13 @@ def aten_trace_backward(grad: TensorType, sizes: INT64) -> TensorType:
     raise NotImplementedError()
 
 
+def aten_transpose(self, dim0: int, dim1: int):
+    # transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)
+
+    # FIXME(justinchuby): onnxscript raises Unsupported expression type
+    return op.Transpose(self, [dim0, dim1])
+
+
 def aten_triangular_solve(
     self: TensorType,
     A: TensorType,

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -150,6 +150,12 @@ def aten_binary_cross_entropy_backward(
     raise NotImplementedError()
 
 
+def aten_celu(self, alpha: float = 1.0):
+    # celu(Tensor self, Scalar alpha=1.0) -> Tensor
+
+    raise NotImplementedError()
+
+
 def aten_col2im(
     self: TensorType,
     output_size: INT64,
@@ -407,12 +413,21 @@ def aten_leaky_relu_backward(
     raise NotImplementedError()
 
 
-def aten_linear(
-    input: TensorType, weight: TensorType, bias: Optional[TensorType] = None
-) -> TensorType:
+def aten_linear(input, weight, bias=None) -> TensorType:
     # linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor
 
-    raise NotImplementedError()
+    # FIXME(justinchuby): Enable the test
+    # INVALID_GRAPH : This is an invalid model.
+    # In Node, ("", OptionalHasElement, "", -1) : () -> ("output0",) ,
+    # Error Node () has input size 0 not in range [min=1, max=1]
+
+    # NOTE: The symbolic function in torch.onnx also uses Gemm in certain cases
+    # Optimizers may consider this path and replace it with Gemm
+    result = op.MatMul(input, weight)
+    if op.OptionalHasElement(bias):
+        bias = op.OptionalGetElement(bias)
+        result = op.Add(result, bias)  # type: ignore[arg-type]
+    return result
 
 
 def aten_log_sigmoid(self: TensorType) -> TensorType: