Update on "extend torch.jit._overload to module methods"

root · root · commit e58f6c6e9fa4 · 2019-08-16T09:07:33.000-07:00
Follow up to #23886, add the same overload api specified in PEP 484 to module methods to reduce the friction of adding method overloads that was brought up in #23266. The usage is: ``` @torch.jit.overload def add(self, y: int) -> int: ... @torch.jit.overload def add(self, y: float) -> float: ... def add(): ... ```
diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
@@ -751,7 +751,7 @@ void initJitScriptBindings(PyObject* module) {
         return script_compile_function(name, new_def, defaults, std::move(rcb));
       });
   m.def(
-      "replace_overloaded_method_decl",
+      "_replace_overloaded_method_decl",
       [](const Decl& overload_decl,
          const Def& implementation_def,
          const std::string& new_name) {
diff --git a/torch/jit/_recursive.py b/torch/jit/_recursive.py
@@ -142,7 +142,7 @@ def recursive_script(mod):
         for overload_fn, name in zip(overload_fns, names):
             torch.jit._check_no_signature(overload_fn)
             over_ast = torch.jit.get_jit_def(overload_fn, self_name="ScriptModule")
-            new_ast = torch._C.replace_overloaded_method_decl(over_ast.decl(), orig_ast, name)
+            new_ast = torch._C._replace_overloaded_method_decl(over_ast.decl(), orig_ast, name)
             _rcb = _jit_internal.createResolutionCallbackFromClosure(orig_fn)
             overload_stubs.append(torch.jit.ScriptMethodStub(_rcb, new_ast, overload_fn))
 
