Add eager ops for unary ops with out. (#12106)

orttraining/orttraining/eager/opgen/opgen/atenops.py

@@ -46,6 +46,57 @@ class GeluGrad(ONNXOp):
 ops = {}
 type_promotion_ops = []
 
+# the following op list is for ops that have a .out version. Often this is the only op needing to be implemented
+# and the regular and inplace(_) version derive from the .out.
+# Example: abs.out is the elementary op and abs and abs_ derive from that.
+# However, that is not alway the case. Example: nonzero has .out and a normal version which is not derived.
+# In this case nonzero will also be in the regular list.
+unary_ops_with_out = [
+    "abs",
+    "acos",
+    "acosh",
+    "asin",
+    "asinh",
+    "atan",
+    "atanh",
+    "ceil",
+    "cos",
+    "cosh",
+    "erf",
+    "exp",
+    "floor",
+    "hardsigmoid",
+    "log",
+    "neg",
+    "nonzero",
+    "reciprocal",
+    "round",
+    "sigmoid",
+    "sign",
+    "sin",
+    "sinh",
+    "sqrt",
+    "tan",
+    "tanh",
+]
+
+# the following ops have explicit inplace(_) needing to be implemented.
+unary_ops_with_inplace = [
+    "relu",
+    "selu",
+]
+
+# the following ops have the regular version needing to be implemented.
+# Note det, isinf, selu - are composite aten ops but have direct onnx ops.
+unary_ops = [
+    "det",
+    "isinf",
+    "isnan",
+    "nonzero",
+    "relu",
+    "selu",
+]
+
 for binary_op, onnx_op in {
     "add": Add("self", Mul("alpha", "other")),
     "sub": Sub("self", Mul("alpha", "other")),
@@ -59,44 +110,14 @@ for binary_op, onnx_op in {
                 ops[f"aten::{binary_op}{variant}.{dtype}"] = deepcopy(onnx_op)
                 type_promotion_ops.append(f"aten::{binary_op}{variant}.{dtype}")
 
-for unary_op in [
-    "acos",
-    "acosh",
-    "asinh",
-    "atanh",
-    "asin",
-    "atan",
-    "ceil",
-    "cos",
-    "cosh",
-    "erf",
-    "exp",
-    "floor",
-    "isnan",
-    "log",
-    "reciprocal",
-    "neg",
-    "round",
-    "relu",
-    "selu",
-    "sigmoid",
-    "sin",
-    "sinh",
-    "sqrt",
-    "tan",
-    "tanh",
-    "nonzero",
-    "sign",
-    "hardsigmoid",
-    "isinf",
-    "det",
-]:
-    aten_name = f"aten::{unary_op}"
-    onnx_op = onnx_ops[unary_op]("self")
-    ops[aten_name] = onnx_op
-    # produce the in-place variant as well for ops that support it
-    if unary_op not in ["isnan", "nonzero", "min", "max", "isinf", "det"]:
-        ops[f"{aten_name}_"] = onnx_op
+for unary_op in unary_ops_with_out:
+    ops[f"aten::{unary_op}.out"] = onnx_ops[unary_op]("self")
+
+for unary_op in unary_ops_with_inplace:
+    ops[f"aten::{unary_op}_"] = onnx_ops[unary_op]("self")
+
+for unary_op in unary_ops:
+    ops[f"aten::{unary_op}"] = onnx_ops[unary_op]("self")
 
 # Notes on Onnx op mapping
 #
@@ -107,7 +128,6 @@ for unary_op in [
 # ---------------------------
 
 hand_implemented = {
-    "aten::abs.out": Abs("self"),
     "aten::empty.memory_format": SignatureOnly(),
     "aten::empty_strided": SignatureOnly(),
     "aten::zero_": SignatureOnly(),

orttraining/orttraining/eager/opgen/opgen/generator.py

@@ -334,6 +334,9 @@ class ORTGen:
 
         # if the torch func has a return ref tensor, out is the last param, and self param is the first input
         # then we need to update and return out.
+        # TODO: make this more general to handle cases where the first param is not self such as
+        # - cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!)
+        # - complex.out(Tensor real, Tensor imag, *, Tensor(a!) out) -> Tensor(a!)
         set_out_tensor = False
         last_param = cpp_func.parameters[-1].member
         if (
@@ -510,6 +513,7 @@ class ORTGen:
                 writer.writeline(f"return {last_param.identifier.value};")
                 return
 
+            # TODO: revisit the hardcoded use of TensorList.
             writer.write(f"at::TensorOptions tensor_options = {first_param.identifier.value}")
             if first_param.parameter_type.desugar().identifier_tokens[0].value == "TensorList":
                 writer.write("[0]")