fix(container): improve 'too many indices' error handling

Reordered index validation checks in `_normalize_idx` to consistently
raise `IndexError` when too many dimensions are provided. Previously,
an index without Ellipses may have skipped this check.

Updated the error message from "array is X-dimensional" to
"container is X-dimensional" for better clarity and consistency.

Adjusted existing tests to expect `IndexError` instead of `RuntimeError`
and added a new test case to explicitly verify this error scenario.

Author: Tim Joseph

src/tensorcontainer/tensor_container.py

@@ -373,32 +373,31 @@ def transform_ellipsis_index(self, shape: torch.Size, idx: tuple) -> tuple:
         Transforms an indexing tuple with an ellipsis into an equivalent one without it.
         ...
         """
-        if Ellipsis not in idx:
-            return idx
-
-        ellipsis_count = 0
-        for item in idx:
-            if item is Ellipsis:
-                ellipsis_count += 1
-        if ellipsis_count > 1:
-            raise IndexError("an index can only have a single ellipsis ('...')")
-
-        ellipsis_pos = idx.index(Ellipsis)
-
         # Count how many items in the index "consume" an axis from the original shape.
         # `None` adds a new axis, so it's not counted.
         num_consuming_indices = sum(
             self.get_number_of_consuming_dims(item) for item in idx
         )
 
         rank = len(shape)
-
         if num_consuming_indices > rank:
             raise IndexError(
-                f"too many indices for array: array is {rank}-dimensional, "
+                f"too many indices for container: container is {rank}-dimensional, "
                 f"but {num_consuming_indices} were indexed"
             )
 
+        if Ellipsis not in idx:
+            return idx
+
+        ellipsis_count = 0
+        for item in idx:
+            if item is Ellipsis:
+                ellipsis_count += 1
+        if ellipsis_count > 1:
+            raise IndexError("an index can only have a single ellipsis ('...')")
+
+        ellipsis_pos = idx.index(Ellipsis)
+
         # Calculate slices needed based on the consuming indices
         num_slices_to_add = rank - num_consuming_indices
 

tests/tensor_dataclass/test_setitem.py

@@ -195,8 +195,8 @@ def test_setitem_broadcast_with_tdc(self, test_name, idx, source_shape):
             "too_many_indices",
             (slice(None), slice(None), 0),
             (20, 5),  # Assign a valid TDC
-            RuntimeError,  # Changed from IndexError to RuntimeError
-            r"Issue with key \.(features|labels) and index \(slice\(None, None, None\), slice\(None, None, None\), 0\) for value of shape torch\.Size\(\[20, 5, 10\]\) and type <class 'torch\.Tensor'> and assignment of shape \(20, 5\)",
+            IndexError,
+            r"too many indices for container: container is \d+-dimensional, but \d+ were indexed",
         ),
     ]
 

tests/tensor_dict/test_getitem.py

@@ -124,3 +124,12 @@ def to_device(obj):
     assert normalize_device(nested["b"].device) == normalize_device(
         torch.device(device)
     )
+
+
+def test_invalid_getitem_raises_error():
+    td = TensorDict({"a": torch.randn(2, 3, 4)}, shape=[2, 3])
+    with pytest.raises(
+        IndexError,
+        match="too many indices for container: container is 2-dimensional, but 3 were indexed",
+    ):
+        td[:, :, 0]