User function variable supports Enum argument (#1491)

Author: yanboliang

test/test_misc.py

@@ -2579,6 +2579,26 @@ def f():
 
         self.assertTrue(torch.allclose(f(), torch.tensor([2.0])))
 
+    def test_user_function_variable_supports_enum_argument(self):
+        class Foo(enum.Enum):
+            FOO = 0
+            BAR = 1
+
+        def gn(x, y=Foo.FOO):
+            if y is Foo.FOO:
+                return x
+            else:
+                return x + 1
+
+        def fn(x):
+            return gn(x)
+
+        x = torch.randn(2, 3)
+        ref = fn(x)
+        opt_fn = torchdynamo.optimize("eager", nopython=True)(fn)
+        res = opt_fn(x)
+        self.assertTrue(torch.allclose(ref, res))
+
 
 class CustomFunc(torch.autograd.Function):
     @staticmethod

torchdynamo/variables/__init__.py

@@ -1,6 +1,7 @@
 from .base import VariableTracker
 from .builtin import BuiltinVariable
 from .constant import ConstantVariable
+from .constant import EnumVariable
 from .dicts import ConstDictVariable
 from .dicts import DataClassVariable
 from .dicts import DefaultDictVariable
@@ -50,6 +51,7 @@
     "ContextWrappingVariable",
     "DataClassVariable",
     "DefaultDictVariable",
+    "EnumVariable",
     "FakeItemVariable",
     "GetAttrVariable",
     "GradModeVariable",

torchdynamo/variables/functions.py

@@ -1,3 +1,4 @@
+import enum
 import functools
 import inspect
 import itertools
@@ -27,6 +28,8 @@ def wrap_bound_arg(val, options):
         return cls([wrap_bound_arg(x, options) for x in val], **options)
     elif variables.ConstantVariable.is_literal(val):
         return variables.ConstantVariable(val, **options)
+    elif isinstance(val, enum.Enum):
+        return variables.EnumVariable(val, **options)
     else:
         assert isinstance(val, VariableTracker), typestr(val)
         return val