Make TS recognize input arg name (#73253)

tugsbayasgalan · cyyever · commit 3cf94ef5f4f1 · 2022-03-03T10:05:28.000+08:00
Summary: Pull Request resolved: pytorch/pytorch#73253 This PR allows TS schema_matching to match input arg with self for aten operators. This is because, operators in their functional form have input as paremeter instead of self. fixes: pytorch/pytorch#71994 Test Plan: Imported from OSS Reviewed By: gmagogsfm Differential Revision: D34427556 Pulled By: tugsbayasgalan fbshipit-source-id: 96c2340d605c59634bf6e37db1db6025d93a933a (cherry picked from commit 45a593d73bc5e6308dd80a4a29afed8e318a0a1c)
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -15771,6 +15771,13 @@ def __init__(self,
 
         torch.jit.script(M(2, 3))
 
+    def test_input_keyword_in_schema(self):
+        def f(x):
+            return torch.ceil(input=x)
+
+        inp = torch.randn(10)
+        self.checkScript(f, (inp, ))
+
     def test_module_method_reassignment(self):
         class Foo(torch.nn.Module):
             def __init__(self):
diff --git a/torch/csrc/jit/frontend/schema_matching.cpp b/torch/csrc/jit/frontend/schema_matching.cpp
@@ -230,10 +230,17 @@ static Value* tryMatchArgument(
 
 c10::optional<size_t> findInputWithName(
     const std::string& name,
-    at::ArrayRef<NamedValue> kwargs) {
+    at::ArrayRef<NamedValue> kwargs,
+    bool is_aten) {
   for (const auto i : c10::irange(kwargs.size())) {
-    if (kwargs[i].name() == name)
+    // TS doesn't understand that the self argument in function
+    // scheams is renamed to input for the functional variant
+    if (is_aten && name == "self" && kwargs[i].name() == "input") {
+      return i;
+    }
+    if (kwargs[i].name() == name) {
       return i;
+    }
   }
   return c10::nullopt;
 }
@@ -342,6 +349,13 @@ static c10::optional<MatchedSchema> tryMatchSchema(
   std::vector<Value*> positional_inputs;
   std::vector<bool> used_kwarg(kwargs.size(), false);
 
+  auto schema_namespace = schema.operator_name().getNamespace();
+  bool is_aten = false;
+  if (schema_namespace.has_value()) {
+    if (schema_namespace.value() == "aten") {
+      is_aten = true;
+    }
+  }
   // if we finish the loop will we have consumed all arguments?
   size_t used_args = 0;
   for (const auto schema_i : c10::irange(schema.arguments().size())) {
@@ -386,7 +400,8 @@ static c10::optional<MatchedSchema> tryMatchSchema(
       // used
       actual_named_value = args[used_args];
       used_args++;
-    } else if (auto kwarg_idx = findInputWithName(arg.name(), kwargs)) {
+    } else if (
+        auto kwarg_idx = findInputWithName(arg.name(), kwargs, is_aten)) {
       const NamedValue& nv = kwargs[*kwarg_idx];
       if (used_kwarg[*kwarg_idx]) {
         if (failure_messages) {
diff --git a/torch/csrc/jit/frontend/schema_matching.h b/torch/csrc/jit/frontend/schema_matching.h
@@ -53,7 +53,8 @@ TORCH_API Value* emitBuiltinCall(
 
 TORCH_API c10::optional<size_t> findInputWithName(
     const std::string& name,
-    at::ArrayRef<NamedValue> kwargs);
+    at::ArrayRef<NamedValue> kwargs,
+    bool is_aten = false);
 
 // applies implicit conversion from value trying to turn it into type
 // concrete_type it succeeds if the return_value->isSubtypeOf(concrete_type)
