Repository: pytorch/tensordict

Issue:
[BUG] Bad Index Behavior after indexing an expanded tensordict
## Describe the bug

Hello tensordict team, I'm unsure where this bug comes from, this happened after I called TensorDict.expand() on the tensordict and indexed it with a 2d tensor.

## To Reproduce

Steps to reproduce the behavior.

Please try to provide a minimal example to reproduce the bug. Error messages and stack traces are also helpful.

Please use the markdown code blocks for both code and stack traces.

```python
import torch
import tensordict

a = torch.arange(9).reshape(3,3)
d_1 = tensordict.TensorDict({
    "a": a,
}, batch_size=(3,3))
d_2 = d_1.unsqueeze(0).expand(1000, 3, 3)
idx = torch.arange(3).unsqueeze(0).expand(3,3) # construct a 2d array to index into d_2
d_2_1 = d_2[idx] # this works
print(d_2_1)
d_2_2 = d_2_1[0:] # this errors

```

```bash
Traceback (most recent call last):
    File "/home/quantumcookie/miniconda3/envs/orangerl/lib/python3.8/site-packages/tensordict/base.py", line 264, in __getitem__
    return self._index_tensordict(index)
  File "/home/quantumcookie/miniconda3/envs/orangerl/lib/python3.8/site-packages/tensordict/_td.py", line 773, in _index_tensordict
    names = self._get_names_idx(index)
  File "/home/quantumcookie/miniconda3/envs/orangerl/lib/python3.8/site-packages/tensordict/_td.py", line 1306, in _get_names_idx
    names = [names[i] if i is not None else None for i in idx_to_take]
  File "/home/quantumcookie/miniconda3/envs/orangerl/lib/python3.8/site-packages/tensordict/_td.py", line 1306, in <listcomp>
    names = [names[i] if i is not None else None for i in idx_to_take]
```

## Expected behavior

This should index completely fine

## System info

torch==2.2.0
tensordict==0.3.0

## Additional context

Add any other context about the problem here.

## Reason and Possible fixes

Unsure

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/tensordict/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Some additional info: 
I installed tensordict==0.2.1 instead and it looks like this code would work fine with that version
Thanks for reporting ! I'll look into it

Resolving patch:
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
--- a/tensordict/_lazy.py
+++ b/tensordict/_lazy.py
@@ -532,10 +532,9 @@ def _split_index(self, index):
             else:
                 if _is_number(idx) and cursor < self.stack_dim:
                     num_single += 1
-                if isinstance(
+                if _is_number(idx) or isinstance(
                     idx,
                     (
-                        int,
                         ftdim.Dim,
                         slice,
                         list,
@@ -560,6 +559,16 @@ def _split_index(self, index):
                             selected_td_idx = range(self.shape[i])
                             split_dim = cursor - num_single
                             mask_loc = i
+                    elif cursor < self.stack_dim:
+                        # we know idx is not a single integer, so it must have
+                        # a dimension. We play with num_single, reducing it
+                        # by the number of dims of idx: if idx has 3 dims, our
+                        # indexed tensor will have 2 more dimensions, going in
+                        # the opposite direction of indexing with a single integer,
+                        # smth[torch.tensor(1)].ndim = smth.ndim-1
+                        # smth[torch.tensor([1])].ndim = smth.ndim
+                        # smth[torch.tensor([[1]])].ndim = smth.ndim+1
+                        num_single -= idx.ndim - 1
                     out.append(idx)
                 else:
                     raise TypeError(f"Invalid index type: {type(idx)}.")
@@ -1562,9 +1571,11 @@ def __getitem__(self, index: IndexType) -> T:
                 return torch.cat(result, cat_dim)
         elif is_nd_tensor:
             new_stack_dim = self.stack_dim - num_single + num_none
-            return LazyStackedTensorDict.lazy_stack(
+            out = LazyStackedTensorDict.lazy_stack(
                 [self[idx] for idx in converted_idx.values()], new_stack_dim
             )
+            out._td_dim_name = self._td_dim_name
+            return out
         else:
             if isinteger:
                 for (
diff --git a/tensordict/_td.py b/tensordict/_td.py
--- a/tensordict/_td.py
+++ b/tensordict/_td.py
@@ -1309,6 +1309,9 @@ def is_boolean(idx):
                         idx_to_take.append(None)
                     elif _is_number(_idx):
                         count += 1
+                    elif isinstance(_idx, (torch.Tensor, np.ndarray)):
+                        idx_to_take.extend([count] * _idx.ndim)
+                        count += 1
                     else:
                         idx_to_take.append(count)
                         count += 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
@@ -2589,6 +2589,29 @@ def test_indexed_properties(self, td_name, device):
         assert td_index.is_shared() is td.is_shared()
         assert td_index.device == td.device
 
+    @pytest.mark.parametrize("npy", [False, True])
+    def test_index_tensor_nd_names(self, td_name, device, npy):
+        td = getattr(self, td_name)(device)
+        names = ("a", "b", "c", "d")
+        try:
+            td.refine_names(*names)
+        except RuntimeError:
+            names = td.names
+        tensor_example = torch.zeros(()).expand(td.shape)
+        assert td.names == list(names)
+        index = torch.tensor([[0, 1, 2], [1, 2, 0], [2, 0, 1]])
+        if npy:
+            index = index.numpy()
+        td_idx = td[:, index]
+        assert tensor_example[:, index].shape == td_idx.shape
+        assert td_idx.names == [names[0], names[1], names[1], *names[2:]]
+        td_idx = td[0, index]
+        assert tensor_example[0, index].shape == td_idx.shape
+        assert td_idx.names == [names[1], names[1], *names[2:]]
+        td_idx = td[..., index, :, :]
+        assert tensor_example[..., index, :, :].shape == td_idx.shape
+        assert td_idx.names == [names[0], names[1], names[1], *names[2:]]
+
     def test_inferred_view_size(self, td_name, device):
         if td_name in ("permute_td", "sub_td2"):
             pytest.skip("view incompatible with stride / permutation")
@@ -5895,12 +5918,12 @@ def test_lazy_indexing(self, pos1, pos2, pos3):
         middle = LazyStackedTensorDict.lazy_stack([inner] * 3, 0)
         outer = LazyStackedTensorDict.lazy_stack([middle] * 2, 0)
         outer_dense = outer.to_tensordict()
-        ref_tensor = torch.zeros(2, 3, 4)
         pos1 = self._idx_list[pos1]
         pos2 = self._idx_list[pos2]
         pos3 = self._idx_list[pos3]
         index = (pos1, pos2, pos3)
         result = outer[index]
+        ref_tensor = torch.zeros(outer.shape)
         assert result.batch_size == ref_tensor[index].shape, index
         assert result.batch_size == outer_dense[index].shape, index