fix dataclass arg repr (#843)

Lightning-AI · Jul 23, 2024 · a16d6ff · a16d6ff
1 parent d27b107
commit a16d6ff
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/thunder/core/codeutils.py b/thunder/core/codeutils.py
@@ -225,10 +225,9 @@ def prettyprint(
         name = _generate_dataclass_class_name(x)
         call_repr = []
         for k, v in x.__dict__.items():
-            try:
-                call_repr.append(f"{k}={v.name}")
-            except:
-                call_repr.append(f"{k}={v}")
+            call_repr.append(
+                f"{k}={prettyprint(v, with_type=False, literals_as_underscores=literals_as_underscores, _quote_markers=False)}"
+            )
         call_repr_str = ",".join(call_repr)
         return m(f"{name}({call_repr_str})")
 

diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py
@@ -2789,11 +2789,12 @@ class TestDataclass:
         s: torch.Tensor
         i: int
         f: float
+        g: tuple
 
     def foo(x):
         # TestDataClass as the output and part of the nested output.
-        return TestDataclass(x, x + 2, x.numel(), x.numel() / 2.0), (
-            TestDataclass(x, x + 2, x.numel(), x.numel() / 2.0),
+        return TestDataclass(x, x + 2, x.numel(), x.numel() / 2.0, (x,)), (
+            TestDataclass(x, x + 2, x.numel(), x.numel() / 2.0, (x)),
             {"x": x, "y": x + 3},
         )
 
@@ -2813,6 +2814,7 @@ def _test_container(actual_container, expected_container):
         torch.testing.assert_close(actual_container.s, expected_container.s)
         torch.testing.assert_close(actual_container.i, expected_container.i)
         torch.testing.assert_close(actual_container.f, expected_container.f)
+        torch.testing.assert_close(actual_container.g[0], expected_container.g[0])
 
     _test_container(actual_container, expected_container)
     _test_container(actual_tuple[0], expected_tuple[0])