Fix 'TensorBox' object has no attribute 'is_input_buffer' (#152980)

jansel · pytorchmergebot · commit 180cbf46f2d2 · 2025-05-09T19:58:32.000Z
Summary: Fix for https://fb.workplace.com/groups/1075192433118967/permalink/1664491270855744/ Test Plan: Used reproducer from D74262030 Differential Revision: D74270090 Pull Request resolved: #152980 Approved by: https://github.com/Skylion007, https://github.com/eellison
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -657,6 +657,12 @@ def maybe_get_name(self) -> Optional[str]:
         except NotImplementedError:
             return None
 
+    def is_input_buffer(self) -> bool:
+        try:
+            return self.get_name() in V.graph.graph_inputs
+        except NotImplementedError:
+            return False
+
     def has_large_inner_fn(self, threshold: Optional[int] = None) -> bool:
         return False
 
@@ -7261,6 +7267,9 @@ def realize_hint(self) -> None:
     def unwrap_view(self) -> IRNode:
         return self.data.unwrap_view()
 
+    def is_input_buffer(self) -> bool:
+        return self.data.is_input_buffer()
+
     def freeze_layout(self) -> None:
         return self.data.freeze_layout()
 
@@ -7380,7 +7389,7 @@ def create(data):  # type: ignore[no-untyped-def]
 
 
 class StorageBox(MutableBox):
-    def is_input_buffer(self):  # type: ignore[no-untyped-def]
+    def is_input_buffer(self) -> bool:
         if isinstance(self.data, (InputBuffer, ReinterpretView)):
             return self.data.get_name() in V.graph.graph_inputs
         return False
