Merge pull request #4 from mctigger/device-fix

Fix TensorContainer.to() device propagation

Author: mctigger

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorcontainer"
-version = "0.6.0"
+version = "0.6.1"
 description = "TensorDict-like functionality for PyTorch with PyTree compatibility and torch.compile support"
 authors = [{name="Tim Joseph", email="tim@mctigger.com"}]
 license = {text = "MIT"}

src/tensorcontainer/tensor_container.py

@@ -437,20 +437,6 @@ def _format_path(cls, path: pytree.KeyPath) -> str:
         return formatted_path
 
     def __repr__(self) -> str:
-        # Infer device for representation if not set (this part is unchanged)
-        device_repr = self.device
-        if device_repr is None:
-            try:
-                # Ensure there are leaves before trying to access device
-                # pytree.tree_leaves can return an empty list
-                leaves = pytree.tree_leaves(self)
-                if leaves:
-                    device_repr = leaves[0].device
-            except IndexError:  # Should not happen if leaves is checked
-                pass
-            except Exception:  # Catch any other pytree or attribute errors
-                pass
-
         # Use a consistent indent of 4 spaces, which is standard
         indent = "    "
 
@@ -480,7 +466,7 @@ def _format_item(key, value):
         return (
             f"{self.__class__.__name__}(\n"
             f"{indent}shape={str(self.shape)},\n"
-            f"{indent}device={device_repr},\n"
+            f"{indent}device={self.device},\n"
             f"{indent}items=\n{textwrap.indent(indented_items, indent)}\n{indent}\n"
             f")"
         )
@@ -622,7 +608,9 @@ def reshape(self: Self, *shape: int) -> Self:
 
     def to(self: Self, *args, **kwargs) -> Self:
         with TensorContainer.unsafe_construction():
-            tc = TensorContainer._tree_map(lambda x: x.to(*args, **kwargs), self)
+            leaves, context = self._pytree_flatten()
+            leaves = [leaf.to(*args, **kwargs) for leaf in leaves]
+            tc = self._pytree_unflatten(leaves, context)
 
         device = self.device