[Feature] Add prepared TensorDict copy_at writer

Author: vmoens

tensordict/base.py

@@ -207,6 +207,112 @@ class _NoDefault(enum.IntEnum):
     Self = Any
 
 
+class _TensorDictCopyAtWriter:
+    def __init__(
+        self,
+        destination: TensorDictBase,
+        dim: int,
+        source: TensorDictBase | None = None,
+    ):
+        self.destination = destination
+        self.dim = _maybe_correct_neg_dim(dim, destination.batch_size)
+        self.keys, self.dest_values = destination._items_list(True, True)
+        self._indexed_dest_values: dict[Any, list[Tensor]] = {}
+        if any(not isinstance(dest, Tensor) for dest in self.dest_values):
+            raise TypeError(
+                "prepare_copy_at_ only supports tensor leaves in the destination "
+                "tensordict."
+            )
+        if source is not None:
+            self._source_values(source)
+
+    def _source_values(self, source: TensorDictBase) -> list[Tensor]:
+        if not _is_tensor_collection(type(source)):
+            raise TypeError(
+                "prepare_copy_at_.copy_ expected a TensorDictBase source, "
+                f"got {type(source)}."
+            )
+        source_keys, source_values = source._items_list(True, True)
+        if source_keys != self.keys:
+            _, source_values = source._items_list(
+                True, True, sorting_keys=self.keys, default=None
+            )
+        if len(source_values) != len(self.dest_values):
+            raise RuntimeError(
+                "The source and destination tensordicts must have matching leaves."
+            )
+        if any(not isinstance(source, Tensor) for source in source_values):
+            raise TypeError(
+                "prepare_copy_at_ only supports tensor leaves in the source "
+                "tensordict."
+            )
+        return source_values
+
+    def _index(self, index: IndexType) -> tuple:
+        idx = [slice(None)] * self.destination.batch_dims
+        idx[self.dim] = index
+        return tuple(idx)
+
+    def _cache_key(self, index: IndexType):
+        try:
+            hash(index)
+        except TypeError:
+            return None
+        return index
+
+    def _get_indexed_dest_values(self, index: IndexType) -> list[Tensor]:
+        cache_key = self._cache_key(index)
+        if cache_key is not None:
+            cached = self._indexed_dest_values.get(cache_key)
+            if cached is not None:
+                return cached
+        idx = self._index(index)
+        indexed_dest_values = []
+        for key, dest in zip(self.keys, self.dest_values):
+            try:
+                dest_indexed = dest[idx]
+            except (IndexError, RuntimeError, TypeError) as err:
+                raise IndexError(
+                    f"Could not index destination leaf {key} with index {idx}."
+                ) from err
+            if (
+                dest_indexed is not dest
+                and getattr(dest_indexed, "_base", None) is None
+            ):
+                raise IndexError(
+                    f"Index {idx} does not produce a writable view for key {key}."
+                )
+            indexed_dest_values.append(dest_indexed)
+        if cache_key is not None:
+            self._indexed_dest_values[cache_key] = indexed_dest_values
+        return indexed_dest_values
+
+    def copy_(
+        self,
+        source: TensorDictBase,
+        index: IndexType,
+        *,
+        non_blocking: bool = False,
+    ) -> TensorDictBase:
+        source_values = self._source_values(source)
+        indexed_dest_values = self._get_indexed_dest_values(index)
+        for key, dest_indexed, source_value in zip(
+            self.keys, indexed_dest_values, source_values
+        ):
+            if dest_indexed.shape != source_value.shape:
+                raise RuntimeError(
+                    f"Shape mismatch for key {key}: indexed destination has shape "
+                    f"{dest_indexed.shape}, source has shape {source_value.shape}."
+                )
+        if _foreach_copy_ is not None:
+            copy_fn = _foreach_copy_compiled if is_compiling() else _foreach_copy_
+            copy_fn(indexed_dest_values, source_values, non_blocking=non_blocking)
+        else:
+            for dest, source_value in zip(indexed_dest_values, source_values):
+                dest.copy_(source_value, non_blocking=non_blocking)
+        return self.destination
+
+
 class _BEST_ATTEMPT_INPLACE:
     def __bool__(self):
         # we use an exception to exit when running `inplace = BEST_ATTEMPT_INPLACE if inplace else False`
@@ -8754,6 +8860,29 @@ def copy_at_(
         """See :obj:`TensorDictBase.update_at_`."""
         return self.update_at_(tensordict, idx, non_blocking=non_blocking)
 
+    def prepare_copy_at_(
+        self, dim: int, source: T | None = None
+    ) -> _TensorDictCopyAtWriter:
+        """Creates a reusable writer for repeated indexed copies.
+
+        Args:
+            dim (int): dimension of ``self`` that will receive each source
+                tensordict.
+            source (TensorDictBase, optional): if provided, validates that the
+                source structure is compatible with ``self``.
+
+        Returns:
+            A lightweight writer with a ``copy_(source, index)`` method.
+
+        Examples:
+            >>> dest = TensorDict({"x": torch.zeros(3, 2)}, batch_size=[3, 2])
+            >>> src = TensorDict({"x": torch.ones(3)}, batch_size=[3])
+            >>> writer = dest.prepare_copy_at_(dim=1, source=src)
+            >>> writer.copy_(src, index=0)
+            >>> assert (dest[:, 0] == src).all()
+        """
+        return _TensorDictCopyAtWriter(self, dim=dim, source=source)
+
     def is_empty(self) -> bool:
         """Checks if the tensordict contains any leaf."""
         for _ in self.keys(True, True):

test/test_tensordict.py

@@ -9290,6 +9290,38 @@ def test_update_at_nested_time_slice_locked(self, td_name, device):
         td.update_at_(td0, (slice(None), 2))
         assert (td[:, 2] == td0).all()
 
+    def test_prepare_copy_at_(self, td_name, device):
+        td = TensorDict(
+            {
+                "a": torch.zeros(4, 3, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.zeros(4, 3, 2, device=device)},
+                    batch_size=[4, 3],
+                    device=device,
+                ),
+            },
+            batch_size=[4, 3],
+            device=device,
+        )
+        td0 = TensorDict(
+            {
+                "a": torch.ones(4, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.ones(4, 2, device=device)},
+                    batch_size=[4],
+                    device=device,
+                ),
+            },
+            batch_size=[4],
+            device=device,
+        )
+        writer = td.prepare_copy_at_(dim=1, source=td0)
+        writer.copy_(td0, index=0)
+        writer.copy_(td0.clone().zero_(), index=2)
+        assert (td[:, 0] == td0).all()
+        assert (td[:, 1] == 0).all()
+        assert (td[:, 2] == 0).all()
+
     # This is needed because update in lazy permute/view etc does not behave correctly when
     # legacy is False. When these classes will be deprecated, we can just remove the decorator
     @set_lazy_legacy(True)