Repository: pytorch/tensordict

Issue:
[Feature Request] Using `pad_sequence(... return_mask=True)` leads to confusing debug sessions
## Motivation

Working on computer vision applications, I often encounter situations where certain object models contain a key called "masks", most commonly referring to something like semantic segmentation outputs.

The function `pad_sequence` has an argument `return_mask` which adds a key called also called "masks". This can lead to some confusing behavior or bugs, especially since this may go unnoticed and could be hard to trace back. 

Currently, I cannot control the key that the padding mask is being written to. In general, there is no clear way to separate/modify keys that TensorDict internally uses and my own user keys.

## Solution

Make the interface with `return_mask=True` return the valid-entry mask as a secondary output, e.g.
```python
def pad_sequence(..., return_mask=False) -> TensorDict: ...  # padded tensordict
def pad_sequence(..., return_mask=True) -> tuple[TensorDict, TensorDict]: ... # padded tensordict, valid masks
```

## Alternatives

1. Propose to have all internally used keys exposed as constants and prepended with an underscore, e.g. export a variable `KEY_MASK: Final = "_mask"`, to clearly differentiate internal keys from user keys. 
2. Propose to pass the key (with default "masks") as an argument, either by a value of `str` to `return_mask` or as a separate parameter.
3. Copy-paste the method to my own codebase and change the key.

## Additional context

I understand that this is a minor/niche request, but since PyTorch is central to many computer vision applications, I hope that the adoption of TensorDict could be helped by avoiding the most common keys that users would be inclined to use.

## Checklist

- [x ] I have checked that there is no similar issue in the repo (**required**)



Discussion:
Since it's a separate key, you could do
```
padded = pad_sequence(td, return_masks=True)
masks = padded.pop("masks")
```

That being said I see your point. The reason we pack it all together is that we can keep the function signature unchanged depending on its arguments. One thing I personally don't enjoy is when changing a keyword argument requires me to adapt the return type. For instance, I find the `max` function confusing in pytorch (`tensor.max()` returns a single value, `tensor.max(1)` returns a named tuple). With tensordict we can avoid that in almost all cases, since we can pack the results together.

@dtsaras what do you think about this?
That makes sense, though the initial reason for raising this was that there are many cases where a tensordict could already have a key called "masks", seeing as this is quite a common key (in vision applications). This would be overwritten and unrecoverable (?) after `pad_sequence` padding.
I personally use it as is and I like that all the data associated are packed in the same TD. I am not a big fan of separating the TD to the padded_TD and the associated masks. We can definitely adjust the function if the community has the need to adjust the masks key though. For example, the user can provide a padding_masks_key which can be set to "padding_masks" (in case "masks" is more likely to write over other data) or "masks" by default. 
> where a tensordict could already have a key called "masks", seeing as this is quite a common key (in vision applications)

What about
```
return_masks=True # sets "masks" key
return_masks="bibbidi-bobbidi-boo" # sets "bibbidi-bobbidi-boo" key
return_masks=False # sets nothing 
```
I guess this works but it certainly isn't pretty
haha yeah
It's the kind of "python isn't strongly typed so let me do whatever I want" solution :) 

I guess most people won't really care about prettyness and most will just use the boolean 

Resolving patch:
diff --git a/tensordict/functional.py b/tensordict/functional.py
--- a/tensordict/functional.py
+++ b/tensordict/functional.py
@@ -84,7 +84,7 @@ def pad_sequence(
     padding_value: float = 0.0,
     out: T | None = None,
     device: DeviceType | None = None,
-    return_mask: bool | None = False,
+    return_mask: bool | NestedKey = False,
 ) -> T:
     """Pads a list of tensordicts in order for them to be stacked together in a contiguous format.
 
@@ -97,7 +97,7 @@ def pad_sequence(
             written.
         device (device compatible type, optional): if provded, the device where the
             TensorDict output will be created.
-        return_mask (bool, optional): if ``True``, a "masks" entry will be returned.
+        return_mask (bool or NestedKey, optional): if ``True``, a "masks" entry will be returned. If ``return_mask`` is a nested key (string or tuple of strings), it will be return the masks and be used as the key for the masks entry.
             It contains a tensordict with the same structure as the stacked tensordict where every entry contains the mask of valid values with size ``torch.Size([stack_len, *new_shape])``,
             where `new_shape[pad_dim] = max_seq_length` and the rest of the `new_shape` matches the previous shape of the contained tensors.
 
@@ -132,6 +132,11 @@ def pad_sequence(
     if not list_of_tensordicts:
         raise RuntimeError("list_of_tensordicts cannot be empty")
 
+    masks_key = "masks"
+    if not isinstance(return_mask, bool):
+        masks_key = unravel_key(return_mask)
+        return_mask = True
+
     # check that all tensordict match
     update_batch_size = True
     max_seq_length = float("-inf")
@@ -144,6 +149,10 @@ def pad_sequence(
 
         for key in keys:
             tensor_shape = td.get(key).shape
+
+            if len(tensor_shape) == 0:
+                raise RuntimeError("Cannot pad scalars")
+            
             pos_pad_dim = pad_dim if pad_dim >= 0 else len(tensor_shape) + pad_dim
 
             # track the maximum sequence length to update batch_size accordingly
@@ -160,7 +169,7 @@ def pad_sequence(
 
             if return_mask:
                 tmp_list_of_tensordicts[-1].set(
-                    ("masks", key),
+                    (masks_key, key),
                     torch.ones(mask_shape, dtype=torch.bool),
                 )
     if return_mask:


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -1085,7 +1085,7 @@ def test_pad(self):
         assert torch.equal(padded_td["a"], expected_a)
         padded_td._check_batch_size()
 
-    @pytest.mark.parametrize("make_mask", [True, False])
+    @pytest.mark.parametrize("make_mask", [True, ("bibbidi", "bobbidi", "boo"), False])
     def test_pad_sequence_pad_dim0(self, make_mask):
         pad_dim = 0
         list_td = [
@@ -1115,22 +1115,25 @@ def test_pad_sequence_pad_dim0(self, make_mask):
             [2, 4, 3]
         )  # check the shape of the padded tensor
         assert torch.all(padded_td["b", "c"][0, 2:, :] == 0)  # check the padding
-        if make_mask:
+        if isinstance(make_mask, str) or make_mask:
+            masks_key = "masks"
+            if isinstance(make_mask, str):
+                masks_key = make_mask
             padded_td_without_masks = pad_sequence(
                 list_td, pad_dim=pad_dim, return_mask=False
             )
-            assert "masks" in padded_td.keys()
+            assert masks_key in padded_td.keys()
             assert set(
                 padded_td_without_masks.keys(include_nested=True, leaves_only=True)
-            ) == set(padded_td["masks"].keys(include_nested=True, leaves_only=True))
-            assert not padded_td["masks", "a"].all()
-            assert padded_td["masks", "a"].ndim == pad_dim + 2
-            assert (padded_td["a"][padded_td["masks", "a"]] != 0).all()
-            assert (padded_td["a"][~padded_td["masks", "a"]] == 0).all()
-            assert not padded_td["masks", "b", "c"].all()
-            assert padded_td["masks", "b", "c"].ndim == pad_dim + 2
-            assert (padded_td["b", "c"][padded_td["masks", "b", "c"]] != 0).all()
-            assert (padded_td["b", "c"][~padded_td["masks", "b", "c"]] == 0).all()
+            ) == set(padded_td[masks_key].keys(include_nested=True, leaves_only=True))
+            assert not padded_td[masks_key, "a"].all()
+            assert padded_td[masks_key, "a"].ndim == pad_dim + 2
+            assert (padded_td["a"][padded_td[masks_key, "a"]] != 0).all()
+            assert (padded_td["a"][~padded_td[masks_key, "a"]] == 0).all()
+            assert not padded_td[masks_key, "b", "c"].all()
+            assert padded_td[masks_key, "b", "c"].ndim == pad_dim + 2
+            assert (padded_td["b", "c"][padded_td[masks_key, "b", "c"]] != 0).all()
+            assert (padded_td["b", "c"][~padded_td[masks_key, "b", "c"]] == 0).all()
         else:
             assert "masks" not in padded_td.keys()
 
@@ -1164,22 +1167,25 @@ def test_pad_sequence_pad_dim1(self, make_mask):
             [2, 6, 7]
         )  # check the shape of the padded tensor
         assert torch.all(padded_td["b", "c"][0, :, 3:] == 0)  # check the padding
-        if make_mask:
+        if isinstance(make_mask, str) or make_mask:
+            masks_key = "masks"
+            if isinstance(make_mask, str):
+                masks_key = make_mask
             padded_td_without_masks = pad_sequence(
                 list_td, pad_dim=pad_dim, return_mask=False
             )
-            assert "masks" in padded_td.keys()
+            assert masks_key in padded_td.keys()
             assert set(
                 padded_td_without_masks.keys(include_nested=True, leaves_only=True)
-            ) == set(padded_td["masks"].keys(include_nested=True, leaves_only=True))
-            assert not padded_td["masks", "a"].all()
-            assert padded_td["masks", "a"].ndim == pad_dim + 2
-            assert (padded_td["a"][padded_td["masks", "a"]] != 0).all()
-            assert (padded_td["a"][~padded_td["masks", "a"]] == 0).all()
-            assert not padded_td["masks", "b", "c"].all()
-            assert padded_td["masks", "b", "c"].ndim == pad_dim + 2
-            assert (padded_td["b", "c"][padded_td["masks", "b", "c"]] != 0).all()
-            assert (padded_td["b", "c"][~padded_td["masks", "b", "c"]] == 0).all()
+            ) == set(padded_td[masks_key].keys(include_nested=True, leaves_only=True))
+            assert not padded_td[masks_key, "a"].all()
+            assert padded_td[masks_key, "a"].ndim == pad_dim + 2
+            assert (padded_td["a"][padded_td[masks_key, "a"]] != 0).all()
+            assert (padded_td["a"][~padded_td[masks_key, "a"]] == 0).all()
+            assert not padded_td[masks_key, "b", "c"].all()
+            assert padded_td[masks_key, "b", "c"].ndim == pad_dim + 2
+            assert (padded_td["b", "c"][padded_td[masks_key, "b", "c"]] != 0).all()
+            assert (padded_td["b", "c"][~padded_td[masks_key, "b", "c"]] == 0).all()
         else:
             assert "masks" not in padded_td.keys()