Update on "[JIT] add support for overloading functions"

root · root · commit 142bd596ead5 · 2019-08-07T14:04:31.000-07:00
This is a series of PRs that will allow us to support adding [padding to conv](#22484) and also reduce the friction of adding method overloads that was brought up in #23266. Support for overloaded functions following the specification in [PEP 484](https://www.python.org/dev/peps/pep-0484/#function-method-overloading). The usage is: ``` @torch.jit.overload def add(x: int, y: int) -> int: ... @torch.jit.overload def add(x: float, y: float) -> float: ... def add: return x + y ``` Follow up PRs: - Add same API for methods - A couple of cleanups for functions: - don't require default params specified on the overload as well - potentially error if invocation could be matched to multiple overloads. now it just chooses the first one, mypy does the same thing currently Differential Revision: [D16694863](https://our.internmc.facebook.com/intern/diff/D16694863)
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12837,12 +12837,12 @@ def test_function_overloads(self):
         # decorators. This is fixed on master but not on version 2.1.1.
         # Next version update remove noqa and add @typing.overload annotation
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def test_simple(x1):  # noqa: F811
             # type: (int) -> int
             pass
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def test_simple(x1):  # noqa: F811
             # type: (float) -> float
             pass
@@ -12856,20 +12856,20 @@ def invoke_function():
         self.checkScript(invoke_function, ())
 
         # testing that the functions are cached
-        compiled_fns_1 = torch.jit.get_overloads(test_simple)
-        compiled_fns_2 = torch.jit.get_overloads(test_simple)
+        compiled_fns_1 = torch.jit._get_overloads(test_simple)
+        compiled_fns_2 = torch.jit._get_overloads(test_simple)
         for a, b in zip(compiled_fns_1, compiled_fns_2):
             self.assertIs(a, b)
 
         # currently we take the default values have to be specified in the
         # overload as well - TODO take them from implementation and apply
         # where the type is valid.
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def identity(x1):  # noqa: F811
             # type: (str) -> str
             pass
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def identity(x1=1.0):  # noqa: F811
             # type: (float) -> float
             pass
@@ -12890,19 +12890,18 @@ def schema_match_failure():
             torch.jit.script(schema_match_failure)
         except Exception as e:
             thrown = True
-            e_msg = str(e)
-            self.assertTrue(r"of type 'str'" in e_msg and r"of type 'float" in e_msg)
+            self.assertTrue(r"of type 'str'" in str(e) and r"of type 'float" in str(e))
         self.assertTrue(thrown)
 
         with self.assertRaisesRegex(Exception, "cannot be directly compiled"):
             torch.jit.script(identity)
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def impl_compile_failure(x, y):  # noqa: F811
             # type: (str, str) -> (str)
             pass
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def impl_compile_failure(x, y):  # noqa: F811
             # type: (int, int) -> (int)
             pass
@@ -12918,19 +12917,22 @@ def test():
             torch.jit.script(test)
 
     def test_function_overloading_isinstance(self):
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def my_conv(x, y):  # noqa: F811
             # type: (float, str) -> (float)
             pass
 
-        @torch.jit.overload  # noqa: F811
+        @torch.jit._overload  # noqa: F811
         def my_conv(x, y=2.0):  # noqa: F811
             # type: (float, float) -> (float)
             pass
 
         def my_conv(x, y=2.0):  # noqa: F811
             if isinstance(y, str):
-                return 4.0 - x
+                if y == "hi":
+                    return 4.0 - x
+                else:
+                    return 5.0 - x
             else:
                 return 2.0 + x
 
diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
@@ -178,6 +178,7 @@ void checkOverloadDecl(const Decl& new_decl, const Decl& old_decl) {
   const auto& new_params = new_decl.params();
   const auto& old_params = old_decl.params();
 
+  // TODO. same number of parameters not strictly necessary.
   TORCH_INTERNAL_ASSERT(
       new_params.size() == old_params.size(),
       "Overload must have same number of parameters\n",
@@ -233,6 +234,27 @@ FunctionSchema getSchemaWithNameAndDefaults(
       schema.is_varret());
 }
 
+static StrongFunctionPtr script_compile_function(
+    const c10::QualifiedName& name,
+    const Def& def,
+    const FunctionDefaults& defaults,
+    ResolutionCallback rcb) {
+  auto cu = get_python_cu();
+  auto defined_functions = cu->define(
+      QualifiedName(name.prefix()),
+      {def},
+      {pythonResolver(std::move(rcb))},
+      nullptr,
+      true);
+  TORCH_INTERNAL_ASSERT(defined_functions.size() == 1);
+  auto& defined = defined_functions[0];
+  defined->setSchema(getSchemaWithNameAndDefaults(
+      def.range(), defined->getSchema(), def.name().name(), defaults));
+  StrongFunctionPtr ret(std::move(cu), defined);
+  didFinishEmitFunction(ret);
+  return ret;
+}
+
 struct VISIBILITY_HIDDEN ModuleSelf : public Self {
   ModuleSelf(const Module& m, py::object& py_m)
       : Self(), module_(m), pyModule_(py_m) {}
@@ -705,20 +727,7 @@ void initJitScriptBindings(PyObject* module) {
         C10_LOG_API_USAGE_ONCE("torch.script.compile");
         const auto name = c10::QualifiedName(qualname);
         TORCH_INTERNAL_ASSERT(name.name() == def.name().name());
-        auto cu = get_python_cu();
-        auto defined_functions = cu->define(
-            QualifiedName(name.prefix()),
-            {def},
-            {pythonResolver(std::move(rcb))},
-            nullptr,
-            true);
-        TORCH_INTERNAL_ASSERT(defined_functions.size() == 1);
-        auto& defined = defined_functions[0];
-        defined->setSchema(getSchemaWithNameAndDefaults(
-            def.range(), defined->getSchema(), def.name().name(), defaults));
-        StrongFunctionPtr ret(std::move(cu), defined);
-        didFinishEmitFunction(ret);
-        return ret;
+        return script_compile_function(name, def, defaults, std::move(rcb));
       });
   m.def(
       "_jit_script_compile_overload",
@@ -728,25 +737,9 @@ void initJitScriptBindings(PyObject* module) {
          ResolutionCallback rcb,
          const FunctionDefaults& defaults) {
         const auto name = c10::QualifiedName(qualname);
-        auto cu = get_python_cu();
         checkOverloadDecl(overload_decl, implementation_def.decl());
         auto new_def = implementation_def.withDecl(overload_decl);
-        auto defined_functions = cu->define(
-            QualifiedName(name.prefix()),
-            {new_def},
-            {pythonResolver(std::move(rcb))},
-            nullptr,
-            true);
-        TORCH_INTERNAL_ASSERT(defined_functions.size() == 1);
-        auto& defined = defined_functions[0];
-        defined->setSchema(getSchemaWithNameAndDefaults(
-            new_def.range(),
-            defined->getSchema(),
-            new_def.name().name(),
-            defaults));
-        StrongFunctionPtr ret(std::move(cu), defined);
-        didFinishEmitFunction(ret);
-        return ret;
+        return script_compile_function(name, new_def, defaults, std::move(rcb));
       });
 
   m.def(
diff --git a/torch/csrc/jit/script/python_sugared_value.cpp b/torch/csrc/jit/script/python_sugared_value.cpp
@@ -573,7 +573,8 @@ std::shared_ptr<SugaredValue> toSugaredValue(
 
   py::bool_ isFunction = py::module::import("inspect").attr("isfunction")(obj);
   if (py::cast<bool>(isFunction)) {
-    auto overloads = py::module::import("torch.jit").attr("get_overloads")(obj);
+    auto overloads =
+        py::module::import("torch.jit").attr("_get_overloads")(obj);
     if (!overloads.is_none()) {
       auto compiled_fns = py::cast<std::vector<StrongFunctionPtr>>(overloads);
       return std::make_shared<OverloadedFunctionValue>(std::move(compiled_fns));
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -1154,7 +1154,7 @@ def forward(self, input):
         _compile_and_register_class(obj, _rcb, qualified_name)
         return obj
     else:
-        check_directly_compile_overloaded(obj)
+        _check_directly_compile_overloaded(obj)
         ast = get_jit_def(obj)
         if _rcb is None:
             _rcb = _gen_rcb(obj, _frames_up)
@@ -2091,7 +2091,7 @@ def _get_script_class(name):
 # qualified name => list[compiled fns]
 _compiled_overloaded_fns = {}
 
-def overload(func):
+def _overload(func):
     qual_name = _qualified_name(func)
     global _overloaded_fns
     fn_overload_list = _overloaded_fns.get(qual_name)
@@ -2104,14 +2104,14 @@ def overload(func):
     fn_overload_list.append((torch.jit.get_jit_def(func).decl(), get_default_args(func)))
     return func
 
-def compile_function_with_overload(qual_name, impl_fn, overload_decl, overload_defaults):
+def _compile_function_with_overload(qual_name, impl_fn, overload_decl, overload_defaults):
     impl_ast = torch.jit.get_jit_def(impl_fn)
     _frames_up = 0
     _rcb = _gen_rcb(impl_fn, _frames_up)
     fn = torch._C._jit_script_compile_overload(qual_name, overload_decl, impl_ast, _rcb, overload_defaults)
     return fn
 
-def get_overloads(obj):
+def _get_overloads(obj):
     # check for cached compiled fns
     qual_name = _qualified_name(obj)
     global _compiled_overloaded_fns
@@ -2130,15 +2130,15 @@ def get_overloads(obj):
     # incompatible with a type of parameter in an overload, and other validation.
     # This is still an internal api so for now use defaults from overload
     for overload_decl, overload_defaults in overloads:
-        compiled_fn = compile_function_with_overload(qual_name, obj, overload_decl, overload_defaults)
+        compiled_fn = _compile_function_with_overload(qual_name, obj, overload_decl, overload_defaults)
         compiled_fns.append(compiled_fn)
 
     # cache compilation, remove information stored to do compilation
     _compiled_overloaded_fns[qual_name] = compiled_fns
     del _overloaded_fns[qual_name]
     return compiled_fns
 
-def check_directly_compile_overloaded(obj):
+def _check_directly_compile_overloaded(obj):
     qual_name = _qualified_name(obj)
     global _compiled_overloaded_fns
     global _overloaded_fns
