Fix pickling dltype namedtuples

dltype/_lib/_core.py

@@ -362,6 +362,7 @@ def _inner_dltyped_namedtuple(cls: type[NT]) -> type[NT]:
 
         # Create a new __new__ method that validates on construction
         original_new = cls.__new__
+        original_module = inspect.getmodule(cls)
 
         def validated_new(cls_inner: type[NT], *args: Any, **kwargs: Any) -> NT:  # noqa: ANN401 (these actually can be any type)
             """A new __new__ method that validates the fields upon construction."""
@@ -385,7 +386,10 @@ def validated_new(cls_inner: type[NT], *args: Any, **kwargs: Any) -> NT:  # noqa
             return instance
 
         # Create the new class with our modified __new__ method
-        return cast("type[NT]", type(cls.__name__, (cls,), {"__new__": validated_new}))
+        new_cls = type(cls.__name__, (cls,), {"__new__": validated_new})
+        new_cls.__module__ = cls.__module__  # Preserve the original module
+        setattr(original_module, cls.__name__, new_cls)  # Update the module's reference to the new class
+        return cast("type[NT]", new_cls)
 
     return _inner_dltyped_namedtuple
 

dltype/tests/dltype_test.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import pickle
 import re
 import sys
 import warnings
@@ -1787,3 +1788,20 @@ def func(arg: type) -> None:
             pass
 
     func(int)
+
+
+def test_pickle_namedtuple() -> None:
+    @dltype.dltyped_namedtuple()
+    class MyNamedTuple(NamedTuple):
+        tensor: Annotated[torch.Tensor, dltype.FloatTensor["b c h w"]]
+        mask: Annotated[torch.Tensor, dltype.IntTensor["b h w"]]
+        other: int
+
+    obj = MyNamedTuple(torch.rand(2, 3, 4, 4), torch.randint(0, 2, (2, 4, 4)), 1)
+    pickled_obj = pickle.dumps(obj)
+    unpickled_obj = pickle.loads(pickled_obj)
+
+    assert isinstance(unpickled_obj, MyNamedTuple)
+    assert torch.allclose(unpickled_obj.tensor, obj.tensor)
+    assert torch.equal(unpickled_obj.mask, obj.mask)
+    assert unpickled_obj.other == obj.other