[Feature] Allow update to increment the number of tds in a lazy stack

ghstack-source-id: 5b670a8
Pull Request resolved: #1220

Author: Vincent Moens

tensordict/_lazy.py

@@ -2727,13 +2727,16 @@ def update(
             isinstance(input_dict_or_td, LazyStackedTensorDict)
             and input_dict_or_td.stack_dim == self.stack_dim
         ):
-            if len(input_dict_or_td.tensordicts) != len(self.tensordicts):
+            tds = list(self.tensordicts)
+            if len(input_dict_or_td.tensordicts) > len(self.tensordicts):
+                tds.extend(
+                    [td.copy() for td in input_dict_or_td.tensordicts[len(tds) :]]
+                )
+            elif len(input_dict_or_td.tensordicts) != len(self.tensordicts):
                 raise ValueError(
                     "cannot update stacked tensordicts with different shapes."
                 )
-            for td_dest, td_source in _zip_strict(
-                self.tensordicts, input_dict_or_td.tensordicts
-            ):
+            for td_dest, td_source in _zip_strict(tds, input_dict_or_td.tensordicts):
                 td_dest.update(
                     td_source,
                     clone=clone,

tensordict/base.py

@@ -6556,6 +6556,11 @@ def update(
 
                 .. seealso:: :meth:`~tensordict.is_leaf_nontensor` and :meth:`~tensordict.default_is_leaf`.
 
+        .. note:: When updating a :class:`~tensordict.LazyStackedTensorDict` with N elements with another
+            :class:`~tensordict.LazyStackedTensorDict` with M elements, with M > N, along the stack dimension,
+            the ``update`` method will append copies of the extra tensordicts to the dest (self) lazy stack.
+            This allows users to rely on ``update`` to increment lazy stacks progressively.
+
         Returns:
             self