Repository: pytorch/rl

Issue:
[Feature Request] Automatic tensordict creator
## Motivation

It would be great to have a helper function that automatically creates a `TensorDict` out of a `dict` with a predefined set of rules.
This should not be the default way of creating TensorDict instances but it should be used in several other helper functions through the library to make the usage of TensorDict-based classes easier.

## Intended usage

The function should work like this:
```python
d = {"a": torch.randn(3, 4, 5), "b": torch.randn(3, 4, 2)}
tensordict = make_tensordict(d)
```
this should return a tensordict with the maximum batch size (in this case `[3, 4]`) and the device (if any) of its content (in this case, "cpu").

It should be possible to parameterize some behaviours, e.g. say how many dimensions should the function look at when determining the batch size, or overriding the device.

This funciton should be implemented in the tensordict `utils.py` file.


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
@@ -4732,3 +4732,41 @@ def _expand_to_match_shape(parent_batch_size, tensor, self_batch_dims, self_devi
 #
 #
 # _register_pytree_node(TensorDict, _flatten_tensordict, _unflatten_tensordict)
+
+
+def make_tensordict(
+    batch_size: Optional[Union[Sequence[int], torch.Size, int]] = None,
+    device: Optional[DEVICE_TYPING] = None,
+    **kwargs,  # source
+) -> TensorDict:
+    """
+    Returns a TensorDict created from the keyword arguments.
+
+    If batch_size is not specified, returns the maximum batch size possible
+
+    Args:
+        **kwargs (TensorDict or torch.Tensor): keyword arguments as data source.
+        batch_size (iterable of int, optional): a batch size for the tensordict.
+        device (torch.device or compatible type, optional): a device for the TensorDict.
+    """
+    if batch_size is None:
+        batch_size = _find_max_batch_size(kwargs)
+    return TensorDict(kwargs, batch_size=batch_size, device=device)
+
+
+def _find_max_batch_size(source: Union[TensorDictBase, dict]) -> list[int]:
+    tensor_data = list(source.values())
+    batch_size = []
+    if not tensor_data:  # when source is empty
+        return batch_size
+    curr_dim = 0
+    while True:
+        if tensor_data[0].dim() > curr_dim:
+            curr_dim_size = tensor_data[0].size(curr_dim)
+        else:
+            return batch_size
+        for tensor in tensor_data[1:]:
+            if tensor.dim() <= curr_dim or tensor.size(curr_dim) != curr_dim_size:
+                return batch_size
+        batch_size.append(curr_dim_size)
+        curr_dim += 1


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -20,6 +20,7 @@
     stack as stack_td,
     pad,
     TensorDictBase,
+    make_tensordict,
 )
 from torchrl.data.tensordict.utils import _getitem_batch_size, convert_ellipsis_to_idx
 
@@ -2565,6 +2566,64 @@ def test_memory_lock(method):
     td.set("b", torch.randn(4, 5), inplace=True)
 
 
+class TestMakeTensorDict:
+    def test_create_tensordict(self):
+        tensordict = make_tensordict(a=torch.zeros(3, 4))
+        assert (tensordict["a"] == torch.zeros(3, 4, 1)).all()
+
+    def test_tensordict_batch_size(self):
+        tensordict = make_tensordict()
+        assert tensordict.batch_size == torch.Size([])
+
+        tensordict = make_tensordict(a=torch.randn(3, 4))
+        assert tensordict.batch_size == torch.Size([3, 4])
+
+        tensordict = make_tensordict(a=torch.randn(3, 4), b=torch.randn(3, 4, 5))
+        assert tensordict.batch_size == torch.Size([3, 4])
+
+        nested_tensordict = make_tensordict(c=tensordict, d=torch.randn(3, 5))  # nested
+        assert nested_tensordict.batch_size == torch.Size([3])
+
+        nested_tensordict = make_tensordict(c=tensordict, d=torch.randn(4, 5))  # nested
+        assert nested_tensordict.batch_size == torch.Size([])
+
+        tensordict = make_tensordict(a=torch.randn(3, 4, 2), b=torch.randn(3, 4, 5))
+        assert tensordict.batch_size == torch.Size([3, 4])
+
+        tensordict = make_tensordict(a=torch.randn(3, 4), b=torch.randn(1))
+        assert tensordict.batch_size == torch.Size([])
+
+        tensordict = make_tensordict(
+            a=torch.randn(3, 4), b=torch.randn(3, 4, 5), batch_size=[3]
+        )
+        assert tensordict.batch_size == torch.Size([3])
+
+        tensordict = make_tensordict(
+            a=torch.randn(3, 4), b=torch.randn(3, 4, 5), batch_size=[]
+        )
+        assert tensordict.batch_size == torch.Size([])
+
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_tensordict_device(self, device):
+        tensordict = make_tensordict(
+            a=torch.randn(3, 4), b=torch.randn(3, 4), device=device
+        )
+        assert tensordict.device == device
+        assert tensordict["a"].device == device
+        assert tensordict["b"].device == device
+
+        tensordict = make_tensordict(
+            a=torch.randn(3, 4, device=device),
+            b=torch.randn(3, 4),
+            c=torch.randn(3, 4, device="cpu"),
+            device=device,
+        )
+        assert tensordict.device == device
+        assert tensordict["a"].device == device
+        assert tensordict["b"].device == device
+        assert tensordict["c"].device == device
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)