Repository: pytorch/tensordict

Issue:
[BUG] Indexing discrepancies with numpy

## Describe the bug

TensorDict indexing behavior differs from numpy. In most cases, the numpy behavior seems to make more sense.

## Steps to reproduce and expected behavior

``` python
import numpy as np
from tensordict import TensorDict

nd = np.random.random((3,4))
nd_5d = np.random.random((5, 3, 4, 6, 7))
td = TensorDict({}, batch_size=(3, 4))
td_5d = TensorDict({}, batch_size=(5, 3, 4, 6, 7))

# None following ellipsis are used as the first dimension instead of the last
nd[1, ..., None].shape  # (4, 1)
td[1, ..., None].shape  # torch.Size([1, 4])

# Same issue
nd[None, 1, ..., None].shape  # (1, 4, 1)
td[None, 1, ..., None].shape  # RuntimeError: Not enough dimensions in TensorDict for index provided.

# Multiple None following ellipsis raise RuntimeError
nd[..., :2, None, None].shape  # (3, 2, 1, 1)
td[..., :2, None, None].shape  # RuntimeError: Not enough dimensions in TensorDict for index provided.

# Tuples within tuples are handled as lists by numpy, while we raise a NotImplementedError. 
# Already documented in https://github.com/pytorch-labs/tensordict/issues/99.
nd[(0, 1), (0, 1)].shape  # (2,)
td[(0, 1), (0, 1)].shape  # NotImplementedError: batch dim cannot be computed for type <class 'tuple'>. 

# Very specific edge case. 
nd_5d[2:, [[[0, 1]]], :3, 0].shape  # (1, 1, 2, 3, 3, 7)
td_5d[2:, [[[0, 1]]], :3, 0].shape  # torch.Size([3, 1, 1, 2, 3, 7])
# However, we get the same output when using a list index rather than an integer
nd_5d[2:, [[[0, 1]]], :3, [0]].shape  # (1, 1, 2, 3, 3, 7)
td_5d[2:, [[[0, 1]]], :3, [0]].shape  # torch.Size([1, 1, 2, 3, 3, 7])
```

## System info

Describe the characteristic of your environment:
 * Describe how the library was installed (pip, source, ...)
 * Python version
 * Versions of any other relevant libraries

```python
import tensordict, numpy, sys, torch
print(tensordict.__version__, numpy.__version__, sys.version, sys.platform, torch.__version__)
```
`0.1.0+c03956a 1.24.2 3.10.8 (main, Nov 24 2022, 08:08:27) [Clang 14.0.6 ] darwin 1.13.1`

## 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/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/tensordict/tensordict.py b/tensordict/tensordict.py
--- a/tensordict/tensordict.py
+++ b/tensordict/tensordict.py
@@ -2693,7 +2693,7 @@ def __getitem__(self, idx: IndexType) -> TensorDictBase:
                 "indexing a tensordict with td.batch_dims==0 is not permitted"
             )
 
-        if isinstance(idx, Number):
+        if _is_number(idx):
             return self._index_tensordict((idx,))
 
         if isinstance(idx, list):
@@ -4407,11 +4407,11 @@ def __init__(
             )
             else tuple(idx)
         )
-        # we msut convert ellipsis into slices
-        idx = self._convert_ellipsis(idx, self._source._batch_size)
-        # idx = self._convert_range(idx)
+        self._batch_size = _getitem_batch_size(self._source.batch_size, idx)
+        if any(item is Ellipsis for item in idx):
+            idx = convert_ellipsis_to_idx(idx, self._source.batch_size)
         self.idx = idx
-        self._batch_size = _getitem_batch_size(self._source.batch_size, self.idx)
+
         if batch_size is not None and batch_size != self.batch_size:
             raise RuntimeError("batch_size does not match self.batch_size.")
 
@@ -5325,89 +5325,94 @@ def __contains__(self, item: IndexType) -> bool:
             return any(item is td for td in self.tensordicts)
         return super().__contains__(item)
 
-    def __getitem__(self, item: IndexType) -> TensorDictBase:
-        if isinstance(item, tuple) and len(item) == 1:
-            item = item[0]
-        if item is Ellipsis or (isinstance(item, tuple) and Ellipsis in item):
-            item = convert_ellipsis_to_idx(item, self.batch_size)
-        if isinstance(item, tuple) and sum(
-            isinstance(_item, str) for _item in item
+    def __getitem__(self, index: IndexType) -> TensorDictBase:
+        if isinstance(index, tuple) and len(index) == 1:
+            index = index[0]
+        if index is Ellipsis or (isinstance(index, tuple) and Ellipsis in index):
+            index = convert_ellipsis_to_idx(index, self.batch_size)
+        if index is None:
+            return self.unsqueeze(0)
+        if isinstance(index, tuple) and sum(
+            isinstance(_item, str) for _item in index
         ) not in [
-            len(item),
+            len(index),
             0,
         ]:
             raise IndexError(_STR_MIXED_INDEX_ERROR)
-        if isinstance(item, (list, range)):
-            item = torch.tensor(item, device=self.device)
-        if isinstance(item, tuple) and any(
-            isinstance(sub_index, (list, range)) for sub_index in item
+        if isinstance(index, (list, range)):
+            index = torch.tensor(index, device=self.device)
+        if isinstance(index, tuple) and any(
+            isinstance(sub_index, (list, range)) for sub_index in index
         ):
-            item = tuple(
+            index = tuple(
                 torch.tensor(sub_index, device=self.device)
                 if isinstance(sub_index, (list, range))
                 else sub_index
-                for sub_index in item
+                for sub_index in index
             )
-        if isinstance(item, str):
-            return self.get(item)
-        elif isinstance(item, tuple) and all(
-            isinstance(sub_item, str) for sub_item in item
+        if isinstance(index, str):
+            return self.get(index)
+        elif isinstance(index, tuple) and all(
+            isinstance(sub_item, str) for sub_item in index
         ):
-            out = self.get(item[0])
-            if len(item) > 1:
+            out = self.get(index[0])
+            if len(index) > 1:
                 if not isinstance(out, TensorDictBase):
                     raise RuntimeError(
                         f"Got a {type(out)} when a TensorDictBase instance was expected."
                     )
-                return out.get(item[1:])
+                return out.get(index[1:])
             else:
                 return out
-        elif isinstance(item, Tensor) and item.dtype == torch.bool:
-            return self.masked_select(item)
-        elif (
-            isinstance(item, (Number,))
-            or (isinstance(item, Tensor) and item.ndimension() == 0)
-        ) and self.stack_dim == 0:
-            return self.tensordicts[item]
-        elif isinstance(item, (Tensor, list)) and self.stack_dim == 0:
+        elif isinstance(index, Tensor) and index.dtype == torch.bool:
+            return self.masked_select(index)
+        elif _is_number(index) and self.stack_dim == 0:
+            return self.tensordicts[index]
+        elif isinstance(index, (Tensor, list)) and self.stack_dim == 0:
             out = LazyStackedTensorDict(
-                *[self.tensordicts[_item] for _item in item],
+                *[self.tensordicts[_item] for _item in index],
                 stack_dim=self.stack_dim,
             )
             return out
-        elif isinstance(item, (Tensor, list)) and self.stack_dim != 0:
-            tds = [tensordict[item] for tensordict in self.tensordicts]
+        elif isinstance(index, (Tensor, list)) and self.stack_dim != 0:
+            tds = [tensordict[index] for tensordict in self.tensordicts]
             dim_drop = self.tensordicts[0].ndim - tds[0].ndim
             out = LazyStackedTensorDict(
                 *tds,
                 stack_dim=self.stack_dim - dim_drop,
             )
             return out
-        elif isinstance(item, slice) and self.stack_dim == 0:
+        elif isinstance(index, slice) and self.stack_dim == 0:
             return LazyStackedTensorDict(
-                *self.tensordicts[item], stack_dim=self.stack_dim
+                *self.tensordicts[index], stack_dim=self.stack_dim
             )
-        elif isinstance(item, slice) and self.stack_dim != 0:
+        elif isinstance(index, slice) and self.stack_dim != 0:
             return LazyStackedTensorDict(
-                *[tensordict[item] for tensordict in self.tensordicts],
+                *[tensordict[index] for tensordict in self.tensordicts],
                 stack_dim=self.stack_dim,
             )
-        elif isinstance(item, (slice, Number)):
+        elif isinstance(index, (slice, Number)):
             new_stack_dim = (
-                self.stack_dim - 1 if isinstance(item, Number) else self.stack_dim
+                self.stack_dim - 1 if isinstance(index, Number) else self.stack_dim
             )
             return LazyStackedTensorDict(
-                *[td[item] for td in self.tensordicts],
+                *[td[index] for td in self.tensordicts],
                 stack_dim=new_stack_dim,
             )
-        elif isinstance(item, tuple):
+        elif isinstance(index, tuple):
+            for i, item in enumerate(index):
+                if item is None:
+                    truncated = tuple(
+                        item if j != i else slice(None) for j, item in enumerate(index)
+                    )
+                    return self.unsqueeze(i).__getitem__(truncated)
             # select sub tensordicts
             _sub_item = tuple(
-                _item for i, _item in enumerate(item) if i != self.stack_dim
+                _item for i, _item in enumerate(index) if i != self.stack_dim
             )
 
-            if self.stack_dim < len(item):
-                idx = item[self.stack_dim]
+            if self.stack_dim < len(index):
+                idx = index[self.stack_dim]
                 if isinstance(idx, (Number, slice)):
                     tensordicts = self.tensordicts[idx]
                 elif isinstance(idx, Tensor):
@@ -5426,14 +5431,24 @@ def __getitem__(self, item: IndexType) -> TensorDictBase:
 
             if len(_sub_item):
                 tensordicts = [td[_sub_item] for td in tensordicts]
+            index_to_stack = []
+            count = 0
+            for item in index:
+                if item is None:
+                    index_to_stack += [item]
+                if count == self.stack_dim:
+                    break
+                count += 1
+                if item is not None:
+                    index_to_stack += [item]
             new_stack_dim = self.stack_dim - sum(
-                [isinstance(_item, Number) for _item in item[: self.stack_dim]]
+                int(_is_number(_item)) - int(_item is None) for _item in index_to_stack
             )
             return torch.stack(list(tensordicts), dim=new_stack_dim)
         else:
             raise NotImplementedError(
                 f"selecting StackedTensorDicts with type "
-                f"{item.__class__.__name__} is not supported yet"
+                f"{index.__class__.__name__} is not supported yet"
             )
 
     def __eq__(self, other):
@@ -6630,3 +6645,11 @@ def _clone_value(value: CompatibleType, recurse: bool) -> CompatibleType:
         return value.clone(recurse=False)
     else:
         return value
+
+
+def _is_number(item):
+    if isinstance(item, Number):
+        return True
+    if isinstance(item, Tensor) and item.ndim == 0:
+        return True
+    return False
diff --git a/tensordict/utils.py b/tensordict/utils.py
--- a/tensordict/utils.py
+++ b/tensordict/utils.py
@@ -75,6 +75,13 @@ def _getitem_batch_size(shape: torch.Size, items: IndexType) -> torch.Size:
 
     Returns:
         Size of the resulting object (tensor or tensordict)
+
+    Examples:
+        >>> idx = (None, ..., None)
+        >>> torch.zeros(4, 3, 2, 1)[idx].shape
+        torch.Size([1, 4, 3, 2, 1, 1])
+        >>> _getitem_batch_size([4, 3, 2, 1], idx)
+        torch.Size([1, 4, 3, 2, 1, 1])
     """
     # let's start with simple cases
     if isinstance(items, tuple) and len(items) == 1:
@@ -96,6 +103,9 @@ def _getitem_batch_size(shape: torch.Size, items: IndexType) -> torch.Size:
     if not isinstance(items, tuple):
         items = (items,)
 
+    if any(item is Ellipsis for item in items):
+        items = convert_ellipsis_to_idx(items, shape)
+
     sanitized_items = []
     for _item in items:
         if isinstance(_item, (list, np.ndarray)):
@@ -122,7 +132,6 @@ def _getitem_batch_size(shape: torch.Size, items: IndexType) -> torch.Size:
     # https://numpy.org/doc/stable/user/basics.indexing.html#integer-array-indexing
     tensor_indices = []
     contiguous, prev = True, None
-
     for i, _item in enumerate(sanitized_items):
         if isinstance(_item, torch.Tensor):
             tensor_indices.append(_item)
@@ -217,8 +226,9 @@ def convert_ellipsis_to_idx(
 
     if idx is Ellipsis:
         idx = (...,)
+
     num_ellipsis = sum(_idx is Ellipsis for _idx in idx)
-    if num_dims < (len(idx) - num_ellipsis):
+    if num_dims < (len(idx) - num_ellipsis - sum(item is None for item in idx)):
         raise RuntimeError("Not enough dimensions in TensorDict for index provided.")
 
     start_pos, after_ellipsis_length = None, 0
@@ -230,6 +240,9 @@ def convert_ellipsis_to_idx(
                 start_pos = i
         if item is not Ellipsis and start_pos is not None:
             after_ellipsis_length += 1
+        if item is None:
+            # unsqueeze
+            num_dims += 1
 
     before_ellipsis_length = start_pos
     ellipsis_length = num_dims - after_ellipsis_length - before_ellipsis_length


Test patch:
diff --git a/test/test_tensordict.py b/test/test_tensordict.py
--- a/test/test_tensordict.py
+++ b/test/test_tensordict.py
@@ -1011,6 +1011,34 @@ 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(
+        "idx",
+        [
+            (..., None),
+            (None, ...),
+            (None,),
+            None,
+            (slice(None), None),
+            (0, None),
+            (None, slice(None), slice(None)),
+            (None, ..., None),
+            (None, 1, ..., None),
+            (1, ..., None),
+            (..., None, 0),
+            ([1], ..., None),
+        ],
+    )
+    def test_index_none(self, td_name, device, idx):
+        td = getattr(self, td_name)(device)
+        tdnone = td[idx]
+        tensor = torch.zeros(td.shape)
+        assert tdnone.shape == tensor[idx].shape, idx
+        if td_name == "td_h5":
+            with pytest.raises(TypeError, match="Selection can't process None"):
+                assert (tdnone.to_tensordict() == td.to_tensordict()[idx]).all()
+            return
+        assert (tdnone.to_tensordict() == td.to_tensordict()[idx]).all()
+
     @pytest.mark.skipif(
         torch.cuda.device_count() == 0, reason="No cuda device detected"
     )
@@ -1457,9 +1485,12 @@ def test_getitem_ellipsis(self, td_name, device, actual_index, expected_index):
 
         actual_td = td[actual_index]
         expected_td = td[expected_index]
+        other_expected_td = td.to_tensordict()[expected_index]
         assert expected_td.shape == _getitem_batch_size(
             td.batch_size, convert_ellipsis_to_idx(actual_index, td.batch_size)
         )
+        assert other_expected_td.shape == actual_td.shape
+        assert_allclose_td(actual_td, other_expected_td)
         assert_allclose_td(actual_td, expected_td)
 
     @pytest.mark.parametrize("actual_index", [..., (..., 0), (0, ...), (0, ..., 0)])