[Doc,BugFix,Feature] Better doc for reductions

ghstack-source-id: 575544e823f29031739dc49561bc2a71125071ca
Pull Request resolved: #1122

tensordict/_td.py

@@ -945,12 +945,18 @@ def _cast_reduction(
                         )
                         for val in agglomerate
                     ]
+                    cat_dim = -1
                     dim = -1
                     keepdim = False
-                agglomerate = torch.cat(agglomerate, dim=-1)
-                return getattr(torch, reduction_name)(
-                    agglomerate, keepdim=keepdim, dim=dim
-                )
+                elif isinstance(dim, tuple):
+                    cat_dim = dim[0]
+                else:
+                    cat_dim = dim
+                agglomerate = torch.cat(agglomerate, dim=cat_dim)
+                kwargs = {}
+                if keepdim is not NO_DEFAULT:
+                    kwargs["keepdim"] = keepdim
+                return getattr(torch, reduction_name)(agglomerate, dim=dim, **kwargs)
 
         # IMPORTANT: do not directly access batch_dims (or any other property)
         # via self.batch_dims otherwise a reference cycle is introduced