[JIT] add support for overloading functions #23886

eellison · 2019-08-06T18:36:56Z

This is a series of PRs that will allow us to support adding padding to conv 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.

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

Stack from ghstack:

[JIT] add support for overloading functions #23886 [JIT] add support for overloading functions
[JIT] metacompile isinstance checks #23885 [JIT] metacompile isinstance checks

Differential Revision: D16694863

ghstack-source-id: efbb867 Pull Request resolved: #23886

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.

ghstack-source-id: 34ded50 Pull Request resolved: #23886

driazati · 2019-08-06T20:49:48Z

test/test_jit.py

        self.assertEqual(out, torch.tensor(6.0))

+    def test_function_overloads(self):
+        # TODO: pyflakes currently does not compose @overload annotation with other


File an issue to update so this doesn't get forgotten

Will do when this lands

test/test_jit.py

torch/csrc/jit/script/init.cpp

torch/jit/__init__.py

driazati · 2019-08-06T21:13:08Z

torch/jit/__init__.py

+        compiled_fns.append(compiled_fn)
+
+    # cache compilation, remove information stored to do compilation
+    _compiled_overloaded_fns[qual_name] = compiled_fns


This looks like a potential memory leak, if the functions this compiles go out of scope, their compiled versions will stick around here. @Chillee found a similar issue recently, I don't remember if he found a fix or not

Discussed in person: this is not a public api yet, and will have very limited usage here so we are going to leave it as is for now.

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

ghstack-source-id: 84e91f6 Pull Request resolved: #23886

torch/jit/__init__.py

Chillee · 2019-08-07T04:04:17Z

Also, can you add a full example somewhere? I'm having trouble getting this to work.

I have

@torch.jit.overload  # noqa: F811
def identity(x1):  # noqa: F811
    # type: (str) -> str
    pass

@torch.jit.overload  # noqa: F811
def identity(x1=1.0):  # noqa: F811
    # type: (float) -> float
    pass

def identity(x1=1.0):  # noqa: F811
    return x1

print(torch.jit.script(identity))

This throws an error:

RuntimeError: Function <function identity at 0x7fba94ef41e0> cannot be directly compiled because it is overloaded. Please wrap it with typed inputs of the signature you wish to compile instead

Specifically, it's not clear to me what "wrapping it with typed inputs" means.

EDIT: I figured it out by looking at the screenshot you initially sent me. I'd suggest rephrasing the error as something like:

RuntimeError: Function <qualname> cannot be directly compiled because it is overloaded. It must be used in the context of a function where `isinstance` checks can statically determine the type.

Perhaps less technically involved. An example might even be a good idea.

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

ghstack-source-id: 715ecba Pull Request resolved: #23886

eellison · 2019-08-07T16:45:37Z

Also, can you add a full example somewhere? I'm having trouble getting this to work.

I have
@torch.jit.overload  # noqa: F811
def identity(x1):  # noqa: F811
    # type: (str) -> str
    pass

@torch.jit.overload  # noqa: F811
def identity(x1=1.0):  # noqa: F811
    # type: (float) -> float
    pass

def identity(x1=1.0):  # noqa: F811
    return x1

print(torch.jit.script(identity))
This throws an error:
RuntimeError: Function <function identity at 0x7fba94ef41e0> cannot be directly compiled because it is overloaded. Please wrap it with typed inputs of the signature you wish to compile instead
Specifically, it's not clear to me what "wrapping it with typed inputs" means.

EDIT: I figured it out by looking at the screenshot you initially sent me. I'd suggest rephrasing the error as something like:
RuntimeError: Function <qualname> cannot be directly compiled because it is overloaded. It must be used in the context of a function where `isinstance` checks can statically determine the type.
Perhaps less technically involved. An example might even be a good idea.

Thanks for the comments. Updated the error message, let me know if you think it's easier to understand. Btw, I suspect this is a very uncommon case (compiling a single function from nn/functional by itself).

driazati

Looks good, remaining stuff is mostly just structural

driazati · 2019-08-07T19:40:48Z

test/test_jit.py

+        def schema_match_failure():
+            return identity((1, 2))
+
+        thrown = False


Why not self.assertRaisesRegex?

Couldn't figure out the regex lol

Would look something like this if you wanna do it with regex (gotta use lookaheads though): https://regexr.com/4irtv

(?=.*of type 'str')(?=.*of type 'float').*

Might not be worth it.

test/test_jit.py

torch/csrc/jit/script/init.cpp

torch/jit/__init__.py

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)

suo · 2019-08-08T02:35:15Z

torch/csrc/jit/script/python_sugared_value.cpp


  py::bool_ isFunction = py::module::import("inspect").attr("isfunction")(obj);
  if (py::cast<bool>(isFunction)) {
+    auto overloads =


Looks like the infrastructure for overloaded functions is only implemented for Python. Do we have a plan for resolving overloaded functions defined from the string frontend (for example, how will we serialize overloaded functions)?

Any overloaded function that needs to be serialized will already be instantiated in one of its overloads, that function gets serialized normally.

suo · 2019-08-08T02:42:47Z

torch/jit/__init__.py

+    qual_name = _qualified_name(obj)
+    global _compiled_overloaded_fns
+    compiled_overloads = _compiled_overloaded_fns.get(qual_name, None)
+    if compiled_overloads is not None:


This doesn't work if we call _get_overloads() then add another overload after. The second overload will not ever get compiled. Do we need a cache here? We don't have anything similar for other recursively compiled functions.

What is the use case you are imagining? This is still an internal only api so I don't think it applies currently but something to think about if we make it public. We don't strictly need a cache it's true.

If we add a new overload after using the overload in the scripted function, it will not be registered. Let's remove the cache, or make the existing caching logic correct. Even if it's unlikely I think it'll be very surprising if we ever run into this.

suo · 2019-08-08T02:50:51Z

torch/csrc/jit/script/init.cpp

          [](Module& self, std::string name, TypePtr type, py::object value) {
-            auto unshaped = unshapedType(type);
-            self.register_attribute(name, unshaped, toIValue(value, type));
+            self.register_attribute(


Why was this changed?

Clang tidy complained, and i accidentally removed the unshaped

facebook-github-bot · 2019-08-08T04:11:36Z

@eellison merged this pull request in 451fc51.

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(): ... ```

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(): ... ``` Differential Revision: [D16921304](https://our.internmc.facebook.com/intern/diff/D16921304)

Summary: Pull Request resolved: #24259 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(): ... ``` Test Plan: Imported from OSS Differential Revision: D16921304 Pulled By: eellison fbshipit-source-id: 784e2f26f7ca9a330a434a603c86b53725c3dc71

add support for overloading functions

90654a6

pytorchbot added oncall: jit Add this issue/PR to JIT oncall triage queue module: pybind Related to our Python bindings / interactions with other Python libraries labels Aug 6, 2019

eellison mentioned this pull request Aug 6, 2019

[JIT] metacompile isinstance checks #23885

Closed

eellison pushed a commit that referenced this pull request Aug 6, 2019

add support for overloading functions

f84c1cc

ghstack-source-id: efbb867 Pull Request resolved: #23886

eellison requested a review from suo August 6, 2019 18:37

eellison changed the title ~~add support for overloading functions~~ [JIT] add support for overloading functions Aug 6, 2019

eellison requested review from Chillee, driazati and zdevito August 6, 2019 18:37

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

397328b

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.

eellison pushed a commit that referenced this pull request Aug 6, 2019

add support for overloading functions

e03adeb

ghstack-source-id: 34ded50 Pull Request resolved: #23886

driazati reviewed Aug 6, 2019

View reviewed changes

eellison requested a review from driazati August 6, 2019 22:56

eellison pushed a commit that referenced this pull request Aug 6, 2019

add support for overloading functions

673f58d

ghstack-source-id: 84e91f6 Pull Request resolved: #23886

Chillee reviewed Aug 7, 2019

View reviewed changes

torch/jit/__init__.py Outdated Show resolved Hide resolved

eellison pushed a commit that referenced this pull request Aug 7, 2019

add support for overloading functions

01b7505

ghstack-source-id: 715ecba Pull Request resolved: #23886

driazati approved these changes Aug 7, 2019

View reviewed changes

facebook-github-bot closed this in 451fc51 Aug 8, 2019

zou3519 deleted the gh/eellison/8/head branch August 8, 2019 02:21

suo reviewed Aug 8, 2019

View reviewed changes

facebook-github-bot added the merged label Aug 8, 2019

eellison mentioned this pull request Aug 13, 2019

extend torch.jit._overload to module methods #24259

Closed

eellison mentioned this pull request Dec 3, 2019

[JIT] polish up overloads on free functions #30356

Closed

mruberry added the Merged label Oct 28, 2020

[JIT] add support for overloading functions #23886

[JIT] add support for overloading functions #23886

Uh oh!

Conversation

eellison commented Aug 6, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Chillee commented Aug 7, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eellison commented Aug 7, 2019

Uh oh!

driazati left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison Aug 8, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison Aug 8, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Aug 8, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

eellison commented Aug 6, 2019 •

edited

Loading

Chillee commented Aug 7, 2019 •

edited

Loading

eellison Aug 8, 2019 •

edited

Loading

eellison Aug 8, 2019 •

edited

Loading