Update (base update)

[ghstack-poisoned]

Author: Vincent Moens

tensordict/nn/probabilistic.py

@@ -15,6 +15,7 @@
 import torch
 
 from tensordict._nestedkey import NestedKey
+from tensordict._td import TensorDict
 from tensordict.base import is_tensor_collection
 from tensordict.nn.common import dispatch, TensorDictModuleBase
 from tensordict.nn.distributions import distributions_maps
@@ -811,6 +812,9 @@ class ProbabilisticTensorDictSequential(TensorDictSequential):
 
             .. warning:: The behaviour of :attr:`return_composite` will change in v0.9
                 and default to True from there on.
+        inplace (bool, optional): if `True`, the input tensordict is modified in-place. If `False`, a new empty
+            :class:`~tensordict.TensorDict` instance is created. If `"empty"`, `input.empty()` is used instead (ie, the
+            output preserves type, device and batch-size). Defaults to `None` (relies on sub-modules).
 
     Raises:
         ValueError: If the input sequence of modules is empty.
@@ -945,6 +949,8 @@ def __init__(
         modules: OrderedDict[str, TensorDictModuleBase | ProbabilisticTensorDictModule],
         partial_tolerant: bool = False,
         return_composite: bool | None = None,
+        *,
+        inplace: bool | None = None,
     ) -> None: ...
 
     @overload
@@ -953,13 +959,16 @@ def __init__(
         modules: List[TensorDictModuleBase | ProbabilisticTensorDictModule],
         partial_tolerant: bool = False,
         return_composite: bool | None = None,
+        *,
+        inplace: bool | None = None,
     ) -> None: ...
 
     def __init__(
         self,
         *modules: TensorDictModuleBase | ProbabilisticTensorDictModule,
         partial_tolerant: bool = False,
         return_composite: bool | None = None,
+        inplace: bool | None = None,
     ) -> None:
         if len(modules) == 0:
             raise ValueError(
@@ -1004,7 +1013,7 @@ def __init__(
         else:
             self.__dict__["_det_part"] = TensorDictSequential(*modules[:-1])
 
-        super().__init__(*modules, partial_tolerant=partial_tolerant)
+        super().__init__(*modules, partial_tolerant=partial_tolerant, inplace=inplace)
         self.return_composite = return_composite
 
     def __getitem__(self, index: int | slice | str) -> TensorDictModuleBase:
@@ -1319,6 +1328,14 @@ def forward(
             tensordict_exec = self._last_module(
                 tensordict_exec, _requires_sample=self._requires_sample
             )
+
+        if self.inplace is True:
+            tensordict_out = tensordict
+        elif self.inplace is False:
+            tensordict_out = TensorDict()
+        elif self.inplace == "empty":
+            tensordict_out = tensordict.empty()
+
         if tensordict_out is not None:
             result = tensordict_out
             result.update(tensordict_exec, keys_to_update=self.out_keys)
@@ -1336,7 +1353,7 @@ def forward(
                     ]
                 else:
                     keys = list(set(self.out_keys + list(tensordict.keys(True, True))))
-                return tensordict.update(result, keys_to_update=keys)
+                return tensordict_out.update(result, keys_to_update=keys)
         return result
 
 

test/test_nn.py

@@ -56,7 +56,7 @@
 )
 
 from torch import distributions, nn
-from torch.distributions import Normal
+from torch.distributions import Categorical, Normal
 from torch.utils._pytree import tree_map
 
 try:
@@ -2127,6 +2127,50 @@ def test_module_buffer():
 
 
 class TestProbabilisticTensorDictModule:
+    @set_composite_lp_aggregate(False)
+    @pytest.mark.parametrize("inplace", [True, False, None])
+    @pytest.mark.parametrize("module_inplace", [True, False])
+    def test_tdprobseq_inplace(self, inplace, module_inplace):
+        model = ProbabilisticTensorDictSequential(
+            TensorDictModule(
+                lambda x: (x + 1, x - 1),
+                in_keys=["input"],
+                out_keys=[("intermediate", "0"), ("intermediate", "1")],
+                inplace=module_inplace,
+            ),
+            TensorDictModule(
+                lambda y0, y1: y0 * y1,
+                in_keys=[("intermediate", "0"), ("intermediate", "1")],
+                out_keys=["output"],
+                inplace=module_inplace,
+            ),
+            ProbabilisticTensorDictModule(
+                in_keys={"logits": "output"},
+                out_keys=["sample"],
+                return_log_prob=True,
+                distribution_class=Categorical,
+            ),
+            inplace=inplace,
+        )
+        input = TensorDict(input=torch.zeros((5,)))
+        output = model(input)
+        assert "sample_log_prob" in output
+        assert "sample" in output
+        if inplace:
+            assert output is input
+            assert "input" in output
+        else:
+            if not module_inplace or inplace is False:
+                # In this case, inplace=False and inplace=None have the same behavior
+                assert output is not input, (module_inplace, inplace)
+                assert "input" not in output, (module_inplace, inplace)
+            else:
+                # In this case, inplace=False and inplace=None have the same behavior
+                assert output is input, (module_inplace, inplace)
+                assert "input" in output, (module_inplace, inplace)
+
+        assert "output" in output
+
     @pytest.mark.parametrize("return_log_prob", [True, False])
     @set_composite_lp_aggregate(False)
     def test_probabilistic_n_samples(self, return_log_prob):