[BugFix] Convert Mapping values for TensorDict fields in tensor-only tensorclasses (#1753)

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>

Author: peterdsharpe

tensordict/tensorclass.py

@@ -19,11 +19,13 @@
 import shutil
 import sys
 import warnings
+from collections.abc import Mapping
 from copy import copy, deepcopy
 from dataclasses import dataclass
 from pathlib import Path
 from textwrap import indent
 from typing import (
+    AbstractSet,
     Any,
     Callable,
     get_args,
@@ -762,19 +764,30 @@ class is created. Without effect if :attr:`obj` is a type.
     """
     from dataclasses import asdict, make_dataclass
 
+    def _field_specs(source, type_hints):
+        return [
+            (
+                field.name,
+                type_hints.get(field.name, field.type),
+                copy(field),
+            )
+            for field in source.__dataclass_fields__.values()
+        ]
+
     if isinstance(obj, type):
         if is_tensorclass(obj):
             return obj
+        type_hints = get_type_hints(obj)
         if not inplace:
             cls = make_dataclass(
                 obj.__name__ + "_tc",
-                fields=obj.__dataclass_fields__,
+                fields=_field_specs(obj, type_hints),
                 bases=obj.__bases__,
             )
         else:
             cls = obj
         clz = _tensorclass(cls, frozen=frozen, shadow=shadow, tensor_only=tensor_only)
-        clz._type_hints = get_type_hints(obj)
+        _set_tensorclass_type_hints(clz, type_hints)
         clz._autocast = autocast
         clz._nocast = nocast
         clz._shadow = shadow
@@ -790,20 +803,37 @@ class is created. Without effect if :attr:`obj` is a type.
             raise TypeError(
                 "tensor_only and autocast or nocast are exclusive features."
             )
+        try:
+            type_hints = get_type_hints(type(obj))
+        except (NameError, TypeError):
+            # Preserve the existing fallback for unresolved annotations.
+            type_hints = None
         clz = _tensorclass(
-            make_dataclass(name, fields=obj.__dataclass_fields__),
+            make_dataclass(
+                name,
+                fields=_field_specs(type(obj), type_hints or {}),
+            ),
             frozen=frozen,
             shadow=shadow,
             tensor_only=tensor_only,
         )
+        if type_hints is not None:
+            _set_tensorclass_type_hints(clz, type_hints)
         clz._autocast = autocast
         clz._nocast = nocast
         clz._shadow = shadow
         clz._frozen = frozen
         clz._tensor_only = tensor_only
     else:
         clz = dest_cls
-    result = clz(**asdict(obj), batch_size=batch_size, device=device)
+    data = asdict(obj)
+    if clz._tensor_only:
+        # ``asdict`` recursively deep-copies values, which would discard the
+        # identity and lock state of TensorDicts. Keep TensorDict-annotated
+        # fields intact and let the tensor-only constructor normalize mappings.
+        for key in clz._tensordict_fields:
+            data[key] = getattr(obj, key)
+    result = clz(**data, batch_size=batch_size, device=device)
     if auto_batch_size:
         if batch_size is not None:
             raise TypeError(
@@ -1442,10 +1472,15 @@ def wrapper(
                 _td_dict = _td._tensordict
                 _non_td = self._non_tensordict
                 _validate = _td._validate_value
+                _tensordict_fields = type(self)._tensordict_fields
                 for key, value in kwargs.items():
                     if value is None:
                         _non_td[key] = None
                     else:
+                        if _tensordict_fields:
+                            value, _ = _convert_mapping_for_field(
+                                key, value, _tensordict_fields
+                            )
                         _td_dict[key] = _validate(
                             value, check_shape=True, non_blocking=False
                         )
@@ -1552,12 +1587,14 @@ def get_parent_locals(cls, localns=localns):
 
     globalns = None
 
+    cls._tensordict_fields = frozenset()
     try:
-        cls._type_hints = get_type_hints(
+        type_hints = get_type_hints(
             cls,
             localns=localns,
             # globalns=globals(),
         )
+        _set_tensorclass_type_hints(cls, type_hints)
         if tensor_only:
 
             def is_tensor_or_optional_tensor(type_hint):
@@ -1632,6 +1669,52 @@ def is_tensor_or_optional_tensor(type_hint):
         cls._type_hints = None
 
 
+def _is_tensordict_annotation(type_hint: Any) -> bool:
+    """Return whether an annotation contains a TensorDictBase type."""
+    origin = get_origin(type_hint)
+    if origin in (Union, UnionType):
+        return any(_is_tensordict_annotation(arg) for arg in get_args(type_hint))
+    if origin is not None:
+        type_hint = origin
+    return isinstance(type_hint, type) and issubclass(type_hint, TensorDictBase)
+
+
+def _set_tensorclass_type_hints(cls: type, type_hints: dict[str, Any]) -> None:
+    """Store resolved hints and cache fields with TensorDict-like annotations."""
+    cls._tensordict_fields = frozenset(
+        key
+        for key, val in type_hints.items()
+        if key in cls.__expected_keys__ and _is_tensordict_annotation(val)
+    )
+    cls._type_hints = type_hints
+
+
+def _normalize_nested_mapping(
+    mapping: Mapping[NestedKey, Any],
+) -> dict[NestedKey, Any] | TensorDictBase:
+    """Materialize nested mappings while preserving TensorDictBase values."""
+    if isinstance(mapping, TensorDictBase):
+        return mapping
+    return {
+        key: _normalize_nested_mapping(value) if isinstance(value, Mapping) else value
+        for key, value in mapping.items()
+    }
+
+
+def _convert_mapping_for_field(
+    key: str, value: Any, tensordict_fields: AbstractSet[str]
+) -> tuple[Any, bool]:
+    """Normalize a Mapping assigned to a TensorDict-annotated field.
+
+    Returns the normalized value and whether the field/value pair required
+    Mapping handling. Existing TensorDict values are returned unchanged.
+    """
+    is_mapping_field = key in tensordict_fields and isinstance(value, Mapping)
+    if is_mapping_field:
+        value = _normalize_nested_mapping(value)
+    return value, is_mapping_field
+
+
 def _from_tensordict(
     cls,
     tensordict: TensorDictBase,
@@ -2023,6 +2106,7 @@ def _setattr_tensor_only(self, key: str, value: Any) -> None:  # noqa: D417
     if value is None:
         self._non_tensordict[key] = None
         return
+    value, _ = _convert_mapping_for_field(key, value, type(self)._tensordict_fields)
     out = self._set_str(key, value, inplace=False, validated=False, ignore_lock=False)
     if out is not self:
         raise RuntimeError(
@@ -2654,6 +2738,13 @@ def set_tensor(
         def _is_castable(datatype):
             return issubclass(datatype, (int, float, np.ndarray))
 
+        if cls._tensor_only:
+            value, is_mapping_field = _convert_mapping_for_field(
+                key, value, cls._tensordict_fields
+            )
+            if is_mapping_field:
+                return set_tensor(value=value)
+
         if cls._autocast:
             type_hints = cls._type_hints
             if type_hints is not None:

test/tensorclass/test_tensorclass.py

@@ -16,6 +16,7 @@
 import re
 import sys
 import weakref
+from collections import UserDict
 from dataclasses import field
 from multiprocessing import Pool
 from pathlib import Path
@@ -35,6 +36,7 @@
     LazyStackedTensorDict,
     MemoryMappedTensor,
     MetaData,
+    NonTensorData,
     set_capture_non_tensor_stack,
     set_list_to_stack,
     tensorclass,
@@ -3619,6 +3621,115 @@ def test_tensor_only_none(self):
         delattr(x, "c")
         assert not hasattr(x, "c")
 
+    @pytest.mark.parametrize("mapping_type", [dict, UserDict])
+    def test_tensor_only_tensordict_mapping(self, mapping_type):
+        @tensorclass(tensor_only=True)
+        class TensorOnlyMapping:
+            data: TensorDict
+
+        value = mapping_type({"tensor": torch.ones(())})
+        tc = TensorOnlyMapping(data=value)
+
+        assert isinstance(tc.data, TensorDict)
+        assert tc.data["tensor"] == 1
+
+        tc.data = mapping_type({"tensor": torch.zeros(())})
+        assert isinstance(tc.data, TensorDict)
+        assert tc.data["tensor"] == 0
+
+        tc.set("data", mapping_type({"tensor": torch.ones(())}))
+        assert isinstance(tc.data, TensorDict)
+        assert tc.data["tensor"] == 1
+
+    def test_tensor_only_tensordict_nested_mapping(self):
+        class TensorOnlyMapping(TensorClass["tensor_only"]):
+            data: TensorDict
+
+        tc = TensorOnlyMapping(
+            data=UserDict({"nested": UserDict({"tensor": torch.ones(())})})
+        )
+
+        assert isinstance(tc.data["nested"], TensorDict)
+        assert tc.data["nested", "tensor"] == 1
+
+    @pytest.mark.parametrize("set_method", ["constructor", "attribute", "set"])
+    def test_tensor_only_preserves_tensordict(self, set_method):
+        class TensorOnlyMapping(TensorClass["tensor_only"]):
+            data: TensorDict
+
+        value = TensorDict({"tensor": torch.ones(())}).lock_()
+        if set_method == "constructor":
+            tc = TensorOnlyMapping(data=value)
+        else:
+            tc = TensorOnlyMapping(data=TensorDict())
+            if set_method == "attribute":
+                tc.data = value
+            else:
+                tc.set("data", value)
+
+        assert tc.data is value
+        assert tc.data.is_locked
+
+    def test_tensor_only_preserves_nested_tensordict(self):
+        class TensorOnlyMapping(TensorClass["tensor_only"]):
+            data: TensorDict
+
+        value = TensorDict({"tensor": torch.ones(())}).lock_()
+        tc = TensorOnlyMapping(data=UserDict({"nested": value}))
+
+        assert tc.data["nested"] is value
+        assert tc.data["nested"].is_locked
+
+    @pytest.mark.parametrize("from_type", [False, True])
+    def test_tensor_only_mapping_from_dataclass(self, from_type):
+        @dataclasses.dataclass
+        class Data:
+            data: TensorDict
+
+        value = UserDict({"tensor": torch.ones(())})
+        if from_type:
+            TensorOnlyData = from_dataclass(Data, tensor_only=True)
+            tc = TensorOnlyData(data=value)
+        else:
+            tc = from_dataclass(Data(data=value), tensor_only=True)
+
+        assert isinstance(tc.data, TensorDict)
+        assert tc.data["tensor"] == 1
+
+    @pytest.mark.parametrize("nested", [False, True])
+    def test_tensor_only_from_dataclass_preserves_tensordict(self, nested):
+        @dataclasses.dataclass
+        class Data:
+            data: TensorDict
+
+        value = TensorDict({"tensor": torch.ones(())}).lock_()
+        data = UserDict({"nested": value}) if nested else value
+
+        tc = from_dataclass(Data(data=data), tensor_only=True)
+        result = tc.data["nested"] if nested else tc.data
+
+        assert result is value
+        assert result.is_locked
+
+    def test_mapping_with_any_annotation_stays_non_tensor(self):
+        class NonTensorMapping(TensorClass):
+            data: Any
+
+        value = UserDict({"metadata": "value"})
+        tc = NonTensorMapping(data=value)
+
+        assert tc.data is value
+
+    def test_tensor_only_non_tensor_mapping_stays_non_tensor(self):
+        class NonTensorMapping(TensorClass["tensor_only"]):
+            data: NonTensorData
+
+        value = UserDict({"metadata": "value"})
+        tc = NonTensorMapping(data=value)
+
+        assert isinstance(tc.data, NonTensorData)
+        assert tc.data.data is value
+
     def test_tensor_only_autocast_nocast(self):
         @tensorclass(tensor_only=True, autocast=False)
         class TensorOnly:
@@ -3696,6 +3807,14 @@ class TensorOnlyGeneric:
             b: TensorDict[str, torch.Tensor]
             c: TensorDict[str, torch.Tensor] | None = None
 
+        tc = TensorOnlyGeneric(
+            a=torch.zeros(()),
+            b=UserDict({"tensor": torch.ones(())}),
+            c=UserDict({"tensor": torch.zeros(())}),
+        )
+        assert isinstance(tc.b, TensorDict)
+        assert isinstance(tc.c, TensorDict)
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()