feat(atenlib): ops 2/n

ghstack-source-id: 23e8ab8
Pull Request resolved: #252

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -3463,16 +3463,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):
@@ -4461,7 +4458,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.Shape(op.Shape(self))
+    if rank == 0 or rank == 1:
+        result = self
+    else:
+        result = op.Transpose(self, perm=[1, 0])
+    return result
 
 
 def aten_t_copy(self: TensorType) -> TensorType:
@@ -4606,6 +4609,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

@@ -407,12 +407,17 @@ 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()
+    result = op.MatMul(input, weight)
+    if op.OptionalHasElement(bias):
+        # FIXME(justinchuby): 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]
+        bias = op.OptionalGetElement(bias)
+        result = op.Add(result, bias)  # type: ignore[arg-type]
+    return result
 
 
 def aten_log_sigmoid(self: TensorType) -> TensorType:

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -162,25 +162,30 @@ def wrapped(fn):
 # Find the names of the OpInfos in torch/testing/_internal/common_methods_invocations.py
 OPINFO_FUNCTION_MAPPING: dict[str, Callable[..., Any]] = {
     "add": core_ops.aten_add,
-    # "clamp": core_ops.aten_clamp,  # TODO(justinchuby): Enable
     "clamp_max": core_ops.aten_clamp_max_tensor,
     "clamp_min": core_ops.aten_clamp_min_tensor,
+    "clamp": core_ops.aten_clamp,
     "gt": core_ops.aten_gt,
     "lt": core_ops.aten_lt,
     "mul": core_ops.aten_mul,
     "nn.functional.elu": nn_ops.aten_elu,
+    "nn.functional.linear": nn_ops.aten_linear,
     "nn.functional.relu6": nn_ops.aten_relu6,
     "nn.functional.selu": core_ops.aten_selu,
     "ones_like": core_ops.aten_ones_like,
+    "ones": core_ops.aten_ones,
     "repeat": core_ops.aten_repeat,
     "round": core_ops.aten_round,
     "sub": core_ops.aten_sub,
+    "t": core_ops.aten_t,
+    # "transpose": core_ops.aten_transpose,  # TODO(justinchuby): Enable when onnxscript errors are fixed
 }
 
 TESTED_OPS = frozenset(OPINFO_FUNCTION_MAPPING)
 
 EXPECTED_SKIPS_OR_FAILS = (
     xfail("add", dtypes=BOOL_TYPES, reason="Add is not defined on bool 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"),
     xfail("gt", dtypes=BOOL_TYPES, reason="Greater is not defined on bool tensors"),
@@ -191,6 +196,10 @@ def wrapped(fn):
         dtypes=dtypes_except(torch.float16, torch.float32),
         reason="ONNX Runtime doesn't support float64 for Elu",
     ),
+    xfail(
+        "nn.functional.linear",
+        reason="ONNX Runtime thinks the graph is invalid",
+    ),
     xfail(
         "nn.functional.relu6",
         dtypes=dtypes_except(torch.float16, torch.float32),
@@ -213,6 +222,7 @@ def wrapped(fn):
         "round", variant_name="decimals_neg_3", reason="The ATen op does not support decimals"
     ),
     xfail("sub", dtypes=BOOL_TYPES, reason="Sub is not defined on bool tensors"),
+    xfail("transpose", reason="Enable when onnxscript errors are fixed"),
 )
 
 
@@ -240,6 +250,10 @@ def wrapped(fn):
 
 OPS_DB = copy.deepcopy(common_methods_invocations.op_db)
 
+ALL_OPS_IN_DB = frozenset(op_info.name for op_info in OPS_DB)
+# Assert all ops in OPINFO_FUNCTION_MAPPING are in the OPS_DB
+assert TESTED_OPS.issubset(ALL_OPS_IN_DB), f"{TESTED_OPS - ALL_OPS_IN_DB} not in OPS_DB"
+
 
 TORCH_TYPE_TO_ONNX = {
     torch.bool: onnx.TensorProto.BOOL,