Use underlying class's __new__ in MutableMappingWrapper.__new__ (Lightning-AI#2514)

Author: t-vi

thunder/core/interpreter.py

@@ -2513,7 +2513,8 @@ def __iter__(self):
 
 class MutMappingWrapperMethods(WrappedValue):
     def __new__(cls, /, *args, **kwds):
-        uvalue = unwrap(cls)()
+        ucls = unwrap(cls)
+        uvalue = ucls.__new__(ucls)
         # todo: for subclasses, better record the call to the constructor
         return wrap(uvalue, provenance=ProvenanceRecord(PseudoInst.NEW, inputs=[cls.provenance]))
 

thunder/tests/test_jit_general.py

@@ -1730,3 +1730,17 @@ def _allocate_and_call_model_in_function():
     ref = _allocate_and_call_model_in_function()
     assert ref() is None
     assert torch.cuda.memory_allocated() == memory_start
+
+
+def test_dataclass_dict():
+    # diffusers model outputs are like this
+    from dataclasses import dataclass
+
+    @dataclass
+    class Foo(dict):
+        musthave: int
+
+    def fn():
+        return Foo(musthave=1)
+
+    assert fn() == thunder.jit(fn)()