Repository: pytorch/tensordict (language: python)

Issue:
[Feature] TensorDict.<reduction>(dim='feature')
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #1122
* __->__ #1121




Discussion / hints:

Fix patch:
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
index 46b636b0e..aeb5a8f23 100644
--- a/tensordict/_lazy.py
+++ b/tensordict/_lazy.py
@@ -2293,15 +2293,44 @@ def _cast_reduction(
         **kwargs,
     ):
         if further_reduce:
-            # It is not very memory-efficient to do this, but it's the easiest to cover all use cases
-            agglomerate = [
-                val.contiguous().flatten()
-                for val in self._values_list(
-                    True, True, is_leaf=_NESTED_TENSORS_AS_LISTS
-                )
-            ]
-            agglomerate = torch.cat(agglomerate, dim=0)
-            return getattr(torch, reduction_name)(agglomerate)
+            if dim is NO_DEFAULT:
+                # It is not very memory-efficient to do this, but it's the easiest to cover all use cases
+                agglomerate = [
+                    val.contiguous().flatten()
+                    for val in self._values_list(
+                        True, True, is_leaf=_NESTED_TENSORS_AS_LISTS
+                    )
+                ]
+                agglomerate = torch.cat(agglomerate, dim=-1)
+                return getattr(torch, reduction_name)(agglomerate)
+            elif dim == "feature":
+
+                def proc_val(val):
+                    val = val.contiguous()
+                    if val.ndim > self.ndim:
+                        val = val.flatten(self.ndim, -1)
+                    else:
+                        val = val.unsqueeze(-1)
+                    return val
+
+                agglomerate = [
+                    proc_val(val)
+                    for val in self.values(
+                        True,
+                        True,
+                    )
+                ]
+                dim = -1
+                cat_dim = -1
+                keepdim = False
+            else:
+                agglomerate = [
+                    val.contiguous().unsqueeze(self.stack_dim)
+                    for val in self.values(True, True)
+                ]
+                cat_dim = self.stack_dim
+            agglomerate = torch.cat(agglomerate, dim=cat_dim)
+            return getattr(torch, reduction_name)(agglomerate, dim=dim, keepdim=keepdim)
 
         try:
             td: TensorDict = self.to_tensordict()
diff --git a/tensordict/_td.py b/tensordict/_td.py
index c50fb437a..deec43b61 100644
--- a/tensordict/_td.py
+++ b/tensordict/_td.py
@@ -936,7 +936,18 @@ def _cast_reduction(
                 agglomerate = list(
                     self._values_list(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS)
                 )
-                agglomerate = torch.cat(agglomerate, dim=0)
+                if dim == "feature":
+                    agglomerate = [
+                        (
+                            val.flatten(self.ndim, -1)
+                            if val.ndim > self.ndim
+                            else val.unsqueeze(-1)
+                        )
+                        for val in agglomerate
+                    ]
+                    dim = -1
+                    keepdim = False
+                agglomerate = torch.cat(agglomerate, dim=-1)
                 return getattr(torch, reduction_name)(
                     agglomerate, keepdim=keepdim, dim=dim
                 )
@@ -963,11 +974,53 @@ def proc_dim(dim, batch_dims, tuple_ok=True):
                 return (batch_dims + dim,)
             return (dim,)
 
-        if dim is not NO_DEFAULT:
+        dim_needs_proc = (dim is not NO_DEFAULT) and (dim not in ("feature",))
+        if dim_needs_proc:
             dim = proc_dim(dim, self.batch_dims, tuple_ok=tuple_ok)
             if not tuple_ok:
                 dim = dim[0]
-        if dim is not NO_DEFAULT or keepdim:
+        if dim in ("feature",):
+            if keepdim:
+                raise TypeError("dim='feature' is incompatible with keepdim=True.")
+
+            ndim = self.ndim
+
+            def reduction(val):
+                if _is_tensor_collection(type(val)):
+                    local_dim = dim
+                else:
+                    if val.ndim > ndim:
+                        val = val.flatten(ndim, -1)
+                    else:
+                        val = val.unsqueeze(-1)
+                    local_dim = -1
+                result = getattr(val, reduction_name)(
+                    dim=local_dim,
+                    **kwargs,
+                )
+                if isinstance(result, tuple):
+                    if values_only:
+                        result = result.values
+                    else:
+                        return TensorDict.from_namedtuple(result)
+                return result
+
+            if self._has_names():
+                names = copy(self.names)
+            else:
+                names = None
+            if not call_on_nested:
+                raise RuntimeError(
+                    f"reduction {reduction_name} must be called with call_on_nested=True when dim='feature'."
+                )
+            return self._fast_apply(
+                reduction,
+                call_on_nested=call_on_nested,
+                device=self.device,
+                names=names,
+            )
+
+        elif dim is not NO_DEFAULT or keepdim:
             names = None
             if self._has_names():
                 names = copy(self.names)
diff --git a/tensordict/base.py b/tensordict/base.py
index e8a9b917a..fec8f9a79 100644
--- a/tensordict/base.py
+++ b/tensordict/base.py
@@ -34,6 +34,7 @@
     Generator,
     Iterator,
     List,
+    Literal,
     Optional,
     OrderedDict,
     overload,
@@ -824,7 +825,7 @@ def cummax(
 
     def mean(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         dtype: torch.dtype | None = None,
@@ -833,11 +834,15 @@ def mean(
         """Returns the mean value of all elements in the input tensordict.
 
         Args:
-            dim (int, tuple of int, optional): if ``None``, returns a dimensionless
+            dim (int, tuple of int, str, optional): if ``None``, returns a dimensionless
                 tensordict containing the mean value of all leaves (if this can be computed).
                 If integer or tuple of integers, `mean` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -862,7 +867,7 @@ def mean(
 
     def nanmean(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         dtype: torch.dtype | None = None,
@@ -876,6 +881,10 @@ def nanmean(
                 If integer or tuple of integers, `mean` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -897,7 +906,7 @@ def nanmean(
 
     def prod(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         dtype: torch.dtype | None = None,
@@ -911,6 +920,10 @@ def prod(
                 If integer or tuple of integers, `prod` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -941,7 +954,7 @@ def prod(
 
     def sum(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         dtype: torch.dtype | None = None,
@@ -955,6 +968,10 @@ def sum(
                 If integer or tuple of integers, `sum` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -976,7 +993,7 @@ def sum(
 
     def nansum(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         dtype: torch.dtype | None = None,
@@ -990,6 +1007,10 @@ def nansum(
                 If integer or tuple of integers, `sum` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -1011,7 +1032,7 @@ def nansum(
 
     def std(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         correction: int = 1,
@@ -1025,6 +1046,10 @@ def std(
                 If integer or tuple of integers, `std` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:
@@ -1045,7 +1070,7 @@ def std(
 
     def var(
         self,
-        dim: int | Tuple[int] = NO_DEFAULT,
+        dim: int | Tuple[int] | Literal["feature"] = NO_DEFAULT,
         keepdim: bool = NO_DEFAULT,
         *,
         correction: int = 1,
@@ -1059,6 +1084,10 @@ def var(
                 If integer or tuple of integers, `var` is called upon the dimension specified if
                 and only if this dimension is compatible with the tensordict
                 shape.
+                Only the `"feature"` string is currently permitted. Using `dim="feature"` will
+                achieve the reduction over all feature dimensions. If `reduce=True`, a tensor of the
+                shape of the TensorDict's batch-size will be returned. Otherwise, a new tensordict
+                with the same structure as ``self`` with reduced feature dimensions will be returned.
             keepdim (bool): whether the output tensor has dim retained or not.
 
         Keyword Args:


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
index d28832b51..6edeb3064 100644
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -2623,6 +2623,56 @@ def test_record_stream(self):
         ):
             td.record_stream(s0)
 
+    @pytest.mark.parametrize(
+        "reduction", ["sum", "nansum", "mean", "nanmean", "std", "var"]
+    )
+    def test_reduction_feature(self, reduction):
+        td = TensorDict(
+            a=torch.ones(3, 4, 5, 6),
+            b=TensorDict(
+                c=torch.ones(3, 4, 5, 6, 7),
+                d=torch.ones(3, 4, 5, 6),
+                e=torch.ones(3, 4, 5) if reduction not in ("std", "var") else None,
+                none=None,
+                batch_size=(3, 4, 5),
+            ),
+            f=torch.ones(3, 4) if reduction not in ("std", "var") else None,
+            none=None,
+            batch_size=(3, 4),
+        )
+        tdr = getattr(td, reduction)(dim="feature")
+        for k, v in td.items(True):
+            other = tdr[k]
+            if isinstance(v, TensorDict):
+                assert v.batch_size == other.batch_size
+            elif is_non_tensor(v):
+                assert v.data is None
+                assert other is None
+            elif isinstance(k, str):
+                assert other.shape == td.shape
+            else:
+                assert other.shape == td[k[:-1]].shape
+
+    @pytest.mark.parametrize(
+        "reduction", ["sum", "nansum", "mean", "nanmean", "std", "var"]
+    )
+    def test_reduction_feature_full(self, reduction):
+        td = TensorDict(
+            a=torch.ones(3, 4, 5, 6),
+            b=TensorDict(
+                c=torch.ones(3, 4, 5, 6, 7),
+                d=torch.ones(3, 4, 5, 6),
+                e=torch.ones(3, 4, 5) if reduction not in ("std", "var") else None,
+                none=None,
+                batch_size=(3, 4, 5),
+            ),
+            f=torch.ones(3, 4) if reduction not in ("std", "var") else None,
+            none=None,
+            batch_size=(3, 4),
+        )
+        reduced = getattr(td, reduction)(dim="feature", reduce=True)
+        assert reduced.shape == (3, 4)
+
     @pytest.mark.parametrize("device", get_available_devices())
     def test_subtensordict_construction(self, device):
         torch.manual_seed(1)
@@ -5529,6 +5579,29 @@ def test_reduction(self, td_name, device, red, tmpdir):
         )
         assert isinstance(getattr(td, red)(reduce=True), torch.Tensor)
 
+    @pytest.mark.parametrize(
+        "red", ("mean", "nanmean", "sum", "nansum", "prod", "std", "var")
+    )
+    def test_reduction_feature(self, td_name, device, red, tmpdir):
+        td = getattr(self, td_name)(device)
+        td = _to_float(td, td_name, tmpdir)
+        if td_name in ("nested_tensorclass", "td_h5"):
+            td = td.apply(
+                lambda x: torch.stack(
+                    [
+                        x,
+                    ]
+                    * 3,
+                    -1,
+                )
+            )
+        tdr = getattr(td, red)(dim="feature")
+        assert tdr.batch_size == td.batch_size
+        assert is_tensor_collection(tdr)
+        tensor = getattr(td, red)(dim="feature", reduce=True)
+        assert tensor.shape == td.batch_size
+        assert isinstance(tensor, torch.Tensor)
+
     @pytest.mark.parametrize("call_del", [True, False])
     def test_remove(self, td_name, device, call_del):
         torch.manual_seed(1)
@@ -8756,6 +8829,31 @@ def test_lazy_stacked_insert(self, dim, index, device):
         with pytest.raises(ValueError, match="Batch sizes in tensordicts differs"):
             lstd.insert(index, TensorDict({"a": torch.ones(17)}, [17], device=device))
 
+    @pytest.mark.parametrize(
+        "reduction", ["sum", "nansum", "mean", "nanmean", "std", "var", "prod"]
+    )
+    def test_reduction_feature_full(self, reduction):
+        td = TensorDict.lazy_stack(
+            [
+                TensorDict.lazy_stack(
+                    [
+                        TensorDict(
+                            a=torch.ones(3, 4),
+                            b=torch.zeros(3, 4, 5),
+                            batch_size=[3, 4],
+                        )
+                        for _ in range(2)
+                    ],
+                    1,
+                )
+                for _ in range(5)
+            ],
+            -1,
+        )
+        assert td.shape == (3, 2, 4, 5)
+        tensor = getattr(td, reduction)(dim="feature", reduce=True)
+        assert tensor.shape == td.shape
+
     @pytest.mark.parametrize("batch_size", [(), (2,), (1, 2)])
     @pytest.mark.parametrize("stack_dim", [0, 1, 2])
     def test_setitem_hetero(self, batch_size, stack_dim):