[dynamo] Trace through dataclasses by removing it from BUILTIN_SKIPLIST (#141294)

StrongerXi · pytorchmergebot · commit 99a0e2b1a1b2 · 2024-11-26T17:05:23.000Z
Fixes #141261. Pull Request resolved: #141294 Approved by: https://github.com/williamwen42, https://github.com/jansel
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -12014,6 +12014,22 @@ def f(x):
         res = f(torch.tensor(1))
         self.assertEqual(torch.tensor(False), res)
 
+    def test_dataclass(self):
+        @dataclasses.dataclass(frozen=True)
+        class Foo:
+            x: int
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def run(x, foo0):
+            if dataclasses.is_dataclass(foo0):
+                foo1 = dataclasses.replace(foo0, **{"x": 1})
+                return x + 1, foo1
+            return x + 2, foo0
+
+        res, foo = run(torch.zeros(1), Foo(0))
+        self.assertTrue(res, torch.ones(1))
+        self.assertEqual(foo.x, 1)
+
 
 class TestTracer(JitTestCase):
     def test_jit_save(self):
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -3155,7 +3155,6 @@ def is_numpy_type_info(obj) -> bool:
     contextlib,
     copy,
     copyreg,
-    dataclasses,
     enum,
     functools,
     importlib,
