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

root · root · commit 98df96dba983 · 2019-08-06T15:48:08.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: int, y: int) -> int: ... 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
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12914,8 +12914,30 @@ def test():
             impl_compile_failure("one", "two")
 
 
-        with self.assertRaisesRegex(Exception, r"# type: (str, str) -> (str"):
-            torch.jit.script(impl_compile_failure)
+        with self.assertRaisesRegex(Exception, "Arguments for call are not valid"):
+            torch.jit.script(test)
+
+    def test_function_overloading_isinstance(self):
+        @torch.jit.overload  # noqa: F811
+        def my_conv(x, y):  # noqa: F811
+            # type: (float, str) -> (float)
+            pass
+
+        @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
+            else:
+                return 2.0 + x
+
+        def test_uses():
+            return my_conv(1.5), my_conv(1.5, "hi"), my_conv(1.5, 5.0)
+
+        self.checkScript(test_uses, ())
 
     @unittest.skipIf(True, "Removing weak script")
     def test_overloading(self):
diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
@@ -180,13 +180,13 @@ void checkOverloadDecl(const Decl& new_decl, const Decl& old_decl) {
 
   TORCH_INTERNAL_ASSERT(
       new_params.size() == old_params.size(),
-      "Overload must have same number of parameters",
+      "Overload must have same number of parameters\n",
       new_decl.range(),
       old_decl.range());
   for (size_t i = 0; i < new_decl.params().size(); ++i) {
     TORCH_INTERNAL_ASSERT(
         new_params[i].ident().name() == old_params[i].ident().name(),
-        "Overload parameters must have the same names",
+        "Overload parameters must have the same names\n",
         new_params[i].ident(),
         old_params[i].ident());
   }
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -1157,22 +1157,28 @@ def forward(self, input):
         check_directly_compile_overloaded(obj)
         ast = get_jit_def(obj)
         if _rcb is None:
-            closure_rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
-            stack_rcb = _jit_internal.createResolutionCallback(_frames_up + 1)
-
-            def _rcb(name):
-                # since type comments aren't captured in the function's closures,
-                # we still need to try to the rcb based on stack frames if the
-                # closure rcb fails
-                result = closure_rcb(name)
-                if result:
-                    return result
-                return stack_rcb(name)
+            _rcb = _gen_rcb(obj, _frames_up)
         fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj))
         # Forward docstrings
         fn.__doc__ = obj.__doc__
         return fn
 
+def _gen_rcb(obj, _frames_up):
+    _frames_up = _frames_up + 1  # for invoking _gen_rcb()
+
+    closure_rcb = _jit_internal.createResolutionCallbackFromClosure(obj)
+    stack_rcb = _jit_internal.createResolutionCallback(_frames_up + 1)
+
+    def _rcb(name):
+        # since type comments aren't captured in the function's closures,
+        # we still need to try to the rcb based on stack frames if the
+        # closure rcb fails
+        result = closure_rcb(name)
+        if result:
+            return result
+        return stack_rcb(name)
+
+    return _rcb
 
 ScriptMethodStub = namedtuple('ScriptMethodStub', ('resolution_callback', 'def_', 'original_method'))
 
@@ -2088,14 +2094,14 @@ def _get_script_class(name):
 def overload(func):
     qual_name = _qualified_name(func)
     global _overloaded_fns
-    li = _overloaded_fns.get(qual_name)
-    if li is None:
-        li = []
-        _overloaded_fns[qual_name] = li
+    fn_overload_list = _overloaded_fns.get(qual_name)
+    if fn_overload_list is None:
+        fn_overload_list = []
+        _overloaded_fns[qual_name] = fn_overload_list
     signature = torch.jit.annotations.get_signature(func)
     if signature is None:
         raise RuntimeError("Must explicitly add type annotations to overloaded functions: {obj}").format(func)
-    li.append((torch.jit.get_jit_def(func).decl(), get_default_args(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):
@@ -2105,15 +2111,8 @@ def compile_function_with_overload(qual_name, impl_fn, overload_decl, overload_d
     # and refactor with above usage
     closure_rcb = _jit_internal.createResolutionCallbackFromClosure(impl_fn)
     stack_rcb = _jit_internal.createResolutionCallback(_frames_up + 1)
-
-    def _rcb(name):
-        # since type comments aren't captured in the function's closures,
-        # we still need to try to the rcb based on stack frames if the
-        # closure rcb fails
-        result = closure_rcb(name)
-        if result:
-            return result
-        return stack_rcb(name)
+    _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
 
