Fix pytorch#121334: graph break on constant method call (pytorch#130158)

Pull Request resolved: pytorch#130158
Approved by: https://github.com/lezcano

Author: rec

test/dynamo/test_functions.py

@@ -330,6 +330,40 @@ def test_methodcall2(a, b):
     def test_methodcall3(a, b):
         return constant3(a, b=1.0) + b
 
+    def test_is_integer(self):
+        @torch.compile(backend="eager", fullgraph=True)
+        def forward(t, m):
+            return 2 * t if m.is_integer() else t
+
+        t = torch.tensor([1])
+        self.assertEqual(forward(t, 1.0).item(), 2)
+        self.assertEqual(forward(t, 1.5).item(), 1)
+
+    @parametrize(
+        "method, num_type",
+        (
+            ("as_integer_ratio", int),
+            ("bit_length", int),
+            ("conjugate", int),
+            ("as_integer_ratio", float),
+            ("conjugate", float),
+            ("hex", float),
+            ("is_integer", float),
+        ),
+    )
+    def test_number_method(self, method, num_type):
+        def forward(t, m):
+            return 2 * t if getattr(m, method)() else t
+
+        wrapped = torch.compile(backend="eager", fullgraph=True)(forward)
+
+        for i in (0, 1, 2.5):
+            m = num_type(i)
+            t = torch.tensor([1])
+            actual = wrapped(t, m)
+            expected = forward(t, m)
+            self.assertEqual(actual, expected)
+
     @make_test
     def test_device_constant(a):
         return a + torch.ones(1, device=torch.device("cpu"))

torch/__init__.py

@@ -511,11 +511,23 @@ def _sympy_(self):
         return self.node.expr
 
     def __hash__(self) -> builtins.int:
+        return hash(self._get_int())
+
+    def as_integer_ratio(self) -> _Tuple[builtins.int, builtins.int]:
+        """Represent this int as an exact integer ratio"""
+        return self._get_int(), 1
+
+    def bit_length(self) -> "SymInt":
+        return SymInt(self.node.wrap_int(self._get_int().bit_length()))
+
+    def conjugate(self) -> "SymInt":
+        return self
+
+    def _get_int(self) -> builtins.int:
         if self.node.is_nested_int():
-            return hash(self.node.nested_int())
+            return self.node.nested_int()
         else:
-            # Force specialization
-            return hash(builtins.int(self))
+            return builtins.int(self)
 
 
 class SymFloat:
@@ -615,18 +627,21 @@ def is_integer(self):
         """Return True if the float is an integer."""
         raise TypeError("type stub not overridden")
 
+    def as_integer_ratio(self) -> _Tuple[builtins.int, builtins.int]:
+        """Represent this float as an exact integer ratio"""
+        return self._get_float().as_integer_ratio()
+
     def __repr__(self):
         return self.node.str()
 
     def _sympy_(self):
         return self.node.expr
 
     def __hash__(self):
-        if self.node.is_constant():
-            return hash(self.node.float_())
-        else:
-            # Force specialization
-            return hash(builtins.float(self))
+        return hash(self._get_float())
+
+    def _get_float(self) -> builtins.float:
+        return self.node.float_() if self.node.is_constant() else builtins.float(self)
 
 
 class SymBool:

torch/_dynamo/variables/constant.py

@@ -155,33 +155,31 @@ def call_method(
         except NotImplementedError:
             return super().call_method(tx, name, args, kwargs)
 
-        def has_arith_binop(num_ty):
-            return (
-                isinstance(self.value, num_ty)
-                and hasattr(operator, name)
-                and len(args) == 1
-                and args[0].is_python_constant()
-            )
-
         if isinstance(self.value, str) and name in str.__dict__.keys():
             method = getattr(self.value, name)
             return ConstantVariable.create(method(*const_args, **const_kwargs))
-        elif has_arith_binop(int) or has_arith_binop(float):
-            op = getattr(operator, name)
-            add_target = const_args[0]
-            if isinstance(add_target, (torch.SymInt, torch.SymFloat)):
-                from .tensor import SymNodeVariable
-
-                # Addition between a non sym and sym makes a sym
-                # sym_num = tx.output.register_attr_or_module(
-                #     add_target, f"sym_shape_{add_target}", source=None
-                # )
-                proxy = tx.output.create_proxy(
-                    "call_function", op, (self.value, add_target), {}
-                )
-                return SymNodeVariable.create(tx, proxy, add_target)
-            return ConstantVariable.create(op(self.value, add_target))
-        elif name == "__len__" and not (args or kwargs):
+        elif isinstance(self.value, (float, int)):
+            if not (args or kwargs):
+                return ConstantVariable.create(getattr(self.value, name)())
+            if (
+                hasattr(operator, name)
+                and len(args) == 1
+                and args[0].is_python_constant()
+            ):
+                add_target = const_args[0]
+                op = getattr(operator, name)
+                if isinstance(
+                    add_target, (torch.SymBool, torch.SymFloat, torch.SymInt)
+                ):
+                    # Addition between a non sym and sym makes a sym
+                    proxy = tx.output.create_proxy(
+                        "call_function", op, (self.value, add_target), {}
+                    )
+                    return SymNodeVariable.create(tx, proxy, add_target)
+                else:
+                    return ConstantVariable.create(op(self.value, add_target))
+
+        if name == "__len__" and not (args or kwargs):
             return ConstantVariable.create(len(self.value))
         elif name == "__contains__" and len(args) == 1 and args[0].is_python_constant():
             assert not kwargs