Enums + jit still not working in PyTorch 2.1 w/ Python 3.11 #108933
Assignees
Labels
oncall: jit
Add this issue/PR to JIT oncall triage queue
triage review
triaged
This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Milestone
Comments
rwightman
changed the title
drisspg
added
oncall: jit
|
FYI to reproduce, w/ a newer (0.9.x) |
|
Original PR that I though should unblock enums for 3.11: #91805 Here is the minimal reproducer for the issue: import torch
from enum import Enum
class Color(int, Enum):
RED = 1
GREEN = 2
def enum_fn(x: Color, y: Color) -> bool:
if x == Color.RED:
return True
return x == y
m = torch.jit.script(enum_fn) |
malfet
added a commit
that referenced
this issue
Sep 20, 2023
In Python-3.11+ typed enums (such as `enum.IntEnum`) retain `__new__`,`__str__` and so on method of the base class via `__init__subclass__()` method (see https://docs.python.org/3/whatsnew/3.11.html#enum ) This change allows typed enums to be scriptable on 3.11, by explicitly marking several `enum.Enum` method to be dropped by jit script Add test that typed enums are jit-scriptable Fixes #108933
malfet
added a commit
that referenced
this issue
Sep 21, 2023
In Python-3.11+ typed enums (such as `enum.IntEnum`) retain `__new__`,`__str__` and so on method of the base class via `__init__subclass__()` method (see https://docs.python.org/3/whatsnew/3.11.html#enum ), i.e. following code ```python import sys import inspect from enum import Enum class IntColor(int, Enum): RED = 1 GREEN = 2 class Color(Enum): RED = 1 GREEN = 2 def get_methods(cls): def predicate(m): if not inspect.isfunction(m) and not inspect.ismethod(m): return False return m.__name__ in cls.__dict__ return inspect.getmembers(cls, predicate=predicate) if __name__ == "__main__": print(sys.version) print(f"IntColor methods {get_methods(IntColor)}") print(f"Color methods {get_methods(Color)}") ``` Returns empty list for both cases for older Python, but on Python-3.11+ it returns list contains of enum constructors and others: ```shell % conda run -n py310 python bar.py 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:41:52) [Clang 15.0.7 ] IntColor methods [] Color methods [] % conda run -n py311 python bar.py 3.11.0 | packaged by conda-forge | (main, Oct 25 2022, 06:21:25) [Clang 14.0.4 ] IntColor methods [('__format__', <function Enum.__format__ at 0x105006ac0>), ('__new__', <function Enum.__new__ at 0x105006660>), ('__repr__', <function Enum.__repr__ at 0x1050068e0>)] Color methods [] ``` This change allows typed enums to be scriptable on 3.11, by explicitly marking several `enum.Enum` method to be dropped by jit script and adds test that typed enums are jit-scriptable. Fixes #108933 Pull Request resolved: #109717 Approved by: https://github.com/atalman, https://github.com/davidberard98 (cherry picked from commit 55685d5)
malfet
added a commit
that referenced
this issue
Sep 21, 2023
In Python-3.11+ typed enums (such as `enum.IntEnum`) retain `__new__`,`__str__` and so on method of the base class via `__init__subclass__()` method (see https://docs.python.org/3/whatsnew/3.11.html#enum ), i.e. following code ```python import sys import inspect from enum import Enum class IntColor(int, Enum): RED = 1 GREEN = 2 class Color(Enum): RED = 1 GREEN = 2 def get_methods(cls): def predicate(m): if not inspect.isfunction(m) and not inspect.ismethod(m): return False return m.__name__ in cls.__dict__ return inspect.getmembers(cls, predicate=predicate) if __name__ == "__main__": print(sys.version) print(f"IntColor methods {get_methods(IntColor)}") print(f"Color methods {get_methods(Color)}") ``` Returns empty list for both cases for older Python, but on Python-3.11+ it returns list contains of enum constructors and others: ```shell % conda run -n py310 python bar.py 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:41:52) [Clang 15.0.7 ] IntColor methods [] Color methods [] % conda run -n py311 python bar.py 3.11.0 | packaged by conda-forge | (main, Oct 25 2022, 06:21:25) [Clang 14.0.4 ] IntColor methods [('__format__', <function Enum.__format__ at 0x105006ac0>), ('__new__', <function Enum.__new__ at 0x105006660>), ('__repr__', <function Enum.__repr__ at 0x1050068e0>)] Color methods [] ``` This change allows typed enums to be scriptable on 3.11, by explicitly marking several `enum.Enum` method to be dropped by jit script and adds test that typed enums are jit-scriptable. Fixes #108933 Cherry-pick of #109717 into release/2.1 branch. Approved by: https://github.com/atalman, https://github.com/davidberard98 (cherry picked from commit 55685d5)
|
Was fixed in lastest 2.1.0 release candidate, tested on my M1 mac by running the following: |
|
@malfet thanks, tested in timm with latest RC and looks good |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
馃悰 Describe the bug
I reported this issue a few months back on torch nightlies where torchscript was having issues with a number of
timmmodels when using Python 3.11. It seemed like there was an awareness of the issues, plan to fix discussing things in slack, but issue persists in the 2.1 RC..It looks like it's a problem with an enum used, so any model that has an enum in timm is failing torchscript tests in python 3.11. No problems with any recent pytorch versions in python < 3.11.
Versions
timm is failing torchscript tests in python 3.11. No problems with any recent pytorch versions in python < 3.11.
cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel
The text was updated successfully, but these errors were encountered: