Repository: pytorch/rl

Issue:
[Feature Request] Make `TensorDict.expand` accept torch.Size/tuple/list arguments
## Motivation

Currently, `TensorDict.expand` only accepts a sequence of integers:
```python
tensordict.expand(a, b, c)
```
but it'd be nice to support `torch.Size` args too:
```python
tensordict.expand(torch.Size([a, b, c]))
```

Ideally this should also work with lists and tuples (i.e. any kind of iterable).



Discussion:

Resolving patch:
diff --git a/torchrl/data/tensordict/tensordict.py b/torchrl/data/tensordict/tensordict.py
--- a/torchrl/data/tensordict/tensordict.py
+++ b/torchrl/data/tensordict/tensordict.py
@@ -670,9 +670,10 @@ def keys(self) -> KeysView:
 
         raise NotImplementedError(f"{self.__class__.__name__}")
 
-    def expand(self, *shape: int) -> TensorDictBase:
+    def expand(self, *shape) -> TensorDictBase:
         """Expands each tensors of the tensordict according to
         `tensor.expand(*shape, *tensor.shape)`
+        Supports iterables to specify the shape
 
         Examples:
             >>> td = TensorDict(source={'a': torch.zeros(3, 4, 5),
@@ -684,6 +685,9 @@ def expand(self, *shape: int) -> TensorDictBase:
         d = dict()
         tensordict_dims = self.batch_dims
 
+        if len(shape) == 1 and isinstance(shape[0], Sequence):
+            shape = tuple(shape[0])
+
         # new shape dim check
         if len(shape) < len(self.shape):
             raise RuntimeError(
@@ -1990,13 +1994,17 @@ def pin_memory(self) -> TensorDictBase:
                     self.set(key, value.pin_memory(), inplace=False)
         return self
 
-    def expand(self, *shape: int) -> TensorDictBase:
+    def expand(self, *shape) -> TensorDictBase:
         """Expands every tensor with `(*shape, *tensor.shape)` and returns the
         same tensordict with new tensors with expanded shapes.
+        Supports iterables to specify the shape.
         """
         d = dict()
         tensordict_dims = self.batch_dims
 
+        if len(shape) == 1 and isinstance(shape[0], Sequence):
+            shape = tuple(shape[0])
+
         # new shape dim check
         if len(shape) < len(self.shape):
             raise RuntimeError(
@@ -3002,7 +3010,10 @@ def select(self, *keys: str, inplace: bool = False) -> TensorDictBase:
             return self
         return self._source.select(*keys)[self.idx]
 
-    def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
+    def expand(self, *shape, inplace: bool = False) -> TensorDictBase:
+        if len(shape) == 1 and isinstance(shape[0], Sequence):
+            shape = tuple(shape[0])
+
         idx = self.idx
         if isinstance(idx, torch.Tensor) and idx.dtype is torch.double:
             # check that idx is not a mask, otherwise throw an error
@@ -3015,7 +3026,6 @@ def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
             idx = idx + (slice(None),) * (self._source.ndimension() - len(idx))
         # now that idx has the same length as the source's number of dims, we can work with it
 
-        # print("shape", shape)
         source_shape = self._source.shape
         num_integer_types = 0
         for i in idx:
@@ -3050,7 +3060,6 @@ def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
                 new_idx.append(_idx)
                 new_source_shape.append(shape[0])
                 shape = shape[1:]
-        # print("source shape", source_shape, "\nnew source shape", new_source_shape, "\nidx", idx, "\nnew_idx", new_idx)
         assert not len(shape)
         new_source = self._source.expand(*new_source_shape)
         new_idx = tuple(new_idx)
@@ -3612,7 +3621,9 @@ def memmap_(self, prefix=None, lock=True) -> TensorDictBase:
         self.is_locked = lock
         return self
 
-    def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
+    def expand(self, *shape, inplace: bool = False) -> TensorDictBase:
+        if len(shape) == 1 and isinstance(shape[0], Sequence):
+            shape = tuple(shape[0])
         stack_dim = len(shape) + self.stack_dim - self.ndimension()
         new_shape_tensordicts = [v for i, v in enumerate(shape) if i != stack_dim]
         tensordicts = [td.expand(*new_shape_tensordicts) for td in self.tensordicts]
@@ -3797,7 +3808,9 @@ def set(self, key: str, value: COMPATIBLE_TYPES, **kwargs) -> TensorDictBase:
         self._save(td)
         return self
 
-    def expand(self, *shape: int, inplace: bool = False) -> TensorDictBase:
+    def expand(self, *shape, inplace: bool = False) -> TensorDictBase:
+        if len(shape) == 1 and isinstance(shape[0], Sequence):
+            shape = tuple(shape[0])
         td = self._load()
         td = td.expand(*shape)
         if inplace:


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -760,10 +760,20 @@ def test_expand(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)
         batch_size = td.batch_size
+        expected_size = torch.Size([3, *batch_size])
+
         new_td = td.expand(3, *batch_size)
-        assert new_td.batch_size == torch.Size([3, *batch_size])
+        assert new_td.batch_size == expected_size
         assert all((_new_td == td).all() for _new_td in new_td)
 
+        new_td_torch_size = td.expand(expected_size)
+        assert new_td_torch_size.batch_size == expected_size
+        assert all((_new_td == td).all() for _new_td in new_td_torch_size)
+
+        new_td_iterable = td.expand([3, *batch_size])
+        assert new_td_iterable.batch_size == expected_size
+        assert all((_new_td == td).all() for _new_td in new_td_iterable)
+
     def test_cast(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)