add support for advanced indexing with less than ndim indexers, ellipsis (pytorch#2144)

Author: killeent

test/test_autograd.py

@@ -555,6 +555,9 @@ def check_index(x, y, idx):
         check_index(x, y, ([slice(None), [2, 3]]))
         check_index(x, y, ([[2, 3], slice(None)]))
 
+        # advanced indexing, with less dim, or ellipsis
+        check_index(x, y, ([0], ))
+
         x = torch.arange(1, 49).view(4, 3, 4)
         y = Variable(x, requires_grad=True)
 
@@ -570,6 +573,14 @@ def check_index(x, y, idx):
         check_index(x, y, (slice(None), [2, 1], slice(None)))
         check_index(x, y, ([2, 1], slice(None), slice(None)))
 
+        # advanced indexing, with less dim, or ellipsis
+        check_index(x, y, ([0], ))
+        check_index(x, y, ([0], slice(None)))
+        check_index(x, y, ([0], Ellipsis))
+        check_index(x, y, ([1, 2], [0, 1]))
+        check_index(x, y, ([1, 2], [0, 1], Ellipsis))
+        check_index(x, y, (Ellipsis, [1, 2], [0, 1]))
+
     def test_indexing_duplicates(self):
         x = torch.arange(1, 17).view(4, 4)
         y = Variable(x, requires_grad=True)
@@ -1458,6 +1469,9 @@ class dont_convert(tuple):
     (Index, (), (torch.rand(S, S, S), dont_convert([slice(None), [0, 3], slice(None)])), 'adv_index_mid'),
     (Index, (), (torch.rand(S, S, S), dont_convert([[0, 3], slice(None), slice(None)])), 'adv_index_beg'),
     (Index, (), (torch.rand(S, S, S), dont_convert([[0, 3], [1, 2], slice(None)])), 'adv_index_comb'),
+    (Index, (), (torch.rand(S, S, S), dont_convert([[0, 3], ])), 'adv_index_sub'),
+    (Index, (), (torch.rand(S, S, S), dont_convert([[0, 3], slice(None)])), 'adv_index_sub_2'),
+    (Index, (), (torch.rand(S, S, S), dont_convert([[0, 3], Ellipsis])), 'adv_index_sub_3'),
     (View, (), (torch.rand(S, S, S), torch.Size([S * S, S]))),
     (Expand, (), ((1, S, 1, S, 1), torch.Size([5, S, 5, S, 5]))),
     (Expand, (), ((S, 1), torch.Size([S, S, S])), 'new_dim'),

test/test_torch.py

@@ -2762,6 +2762,16 @@ def ri(indices):
         self.assertEqual(strided[rows, columns],
                          torch.Tensor([[4, 6], [2, 3]]))
 
+        # Tests using less than the number of dims, and ellipsis
+
+        # reference is 1 2
+        #              3 4
+        #              5 6
+        reference = conv_fn(consec((3, 2)))
+        self.assertEqual(reference[ri([0, 2]), ], torch.Tensor([[1, 2], [5, 6]]))
+        self.assertEqual(reference[ri([1]), ...], torch.Tensor([[3, 4]]))
+        self.assertEqual(reference[..., ri([1])], torch.Tensor([[2], [4], [6]]))
+
         if TEST_NUMPY:
             # we use numpy to compare against, to verify that our advanced
             # indexing semantics are the same, and also for ease of test
@@ -2864,6 +2874,19 @@ def get_set_tensor(indexed, indexer):
                 [[[0, 1], [2, 3]], [[0]], slice(None)],
                 [[[2, 1]], [[0, 3], [4, 4]], slice(None)],
                 [[[2]], [[0, 3], [4, 1]], slice(None)],
+
+                # less dim, ellipsis
+                [[0, 2], ],
+                [[0, 2], slice(None)],
+                [[0, 2], Ellipsis],
+                [[0, 2], slice(None), Ellipsis],
+                [[0, 2], Ellipsis, slice(None)],
+                [[0, 2], [1, 3]],
+                [[0, 2], [1, 3], Ellipsis],
+                [Ellipsis, [1, 3], [2, 3]],
+                [Ellipsis, [2, 3, 4]],
+                [Ellipsis, slice(None), [2, 3, 4]],
+                [slice(None), Ellipsis, [2, 3, 4]],
             ]
 
             for indexer in indices_to_test:
@@ -2917,6 +2940,25 @@ def get_set_tensor(indexed, indexer):
                 [[0], [4], [1, 3, 4], slice(None)],
                 [[1], [0, 2, 3], [1], slice(None)],
                 [[[1, 2], [1, 2]], [[0, 1], [2, 3]], [[2, 3], [3, 5]], slice(None)],
+
+                # less dim, ellipsis
+                [Ellipsis, [0, 3, 4]],
+                [Ellipsis, slice(None), [0, 3, 4]],
+                [Ellipsis, slice(None), slice(None), [0, 3, 4]],
+                [slice(None), Ellipsis, [0, 3, 4]],
+                [slice(None), slice(None), Ellipsis, [0, 3, 4]],
+                [slice(None), [0, 2, 3], [1, 3, 4]],
+                [slice(None), [0, 2, 3], [1, 3, 4], Ellipsis],
+                [Ellipsis, [0, 2, 3], [1, 3, 4], slice(None)],
+                [[0], [1, 2, 4]],
+                [[0], [1, 2, 4], slice(None)],
+                [[0], [1, 2, 4], Ellipsis],
+                [[0], [1, 2, 4], Ellipsis, slice(None)],
+                [[1], ],
+                [[0, 2, 1], [3], [4]],
+                [[0, 2, 1], [3], [4], slice(None)],
+                [[0, 2, 1], [3], [4], Ellipsis],
+                [Ellipsis, [0, 2, 1], [3], [4]],
             ]
 
             for indexer in indices_to_test:

torch/csrc/generic/Tensor.cpp

@@ -531,29 +531,31 @@ static bool THPTensor_(_checkAdvancedIndexing)(THPTensor *indexed, PyObject *arg
   //
   // 1. "Basic Integer Array Indexing" the integer-array indexing strategy
   // where we have ndim sequence/LongTensor arguments
-  // 2. Combining Advanced Indexing with ":", with the limitation that
+  // 2. Combining Advanced Indexing with ":", or "..." , with the limitation that
   // the advanced indexing dimensions must be adjacent, i.e.:
   //
   // x[:, :, [1,2], [3,4], :] --> valid
+  // x[[1,2], [3,4]] --> valid
+  // x[[1,2], [3,4], ...] --> valid
   // x[:, [1,2], :, [3,4], :] --> not valid
 
   // Verification, Step #1 -- ndim sequencers
   if (THPTensor_(_checkBasicIntegerArrayIndexing)(indexed, arg)) return true;
 
   // Verification, Step #2 -- at least one sequencer, all the rest are
-  // ':', can be less than ndim indexers, all sequencers adjacent
+  // ':' and/or a single '...', can be less than ndim indexers, all sequencers
+  // adjacent
 
   long ndim = THTensor_(nDimension)(LIBRARY_STATE indexed->cdata);
-  // TODO: should this be == ndim? --> for now, yes, but to support
-  // other things, no
-  if (PySequence_Check(arg) && PySequence_Size(arg) == ndim) {
+  if (PySequence_Check(arg) && PySequence_Size(arg) <= ndim) {
     THPObjectPtr fast = THPObjectPtr(PySequence_Fast(arg, NULL));
 
     bool sequenceFound = false;
-    bool nonColonFound = false;
+    bool nonColonEllipsisFound = false;
+    bool ellipsisFound = false;
     Py_ssize_t lastSeqDim = -1;
 
-    for (Py_ssize_t i = 0; i < ndim; ++i) {
+    for (Py_ssize_t i = 0; i < PySequence_Fast_GET_SIZE(fast.get()); ++i) {
       PyObject *item = PySequence_Fast_GET_ITEM(fast.get(), i);
       if (THPIndexTensor_Check(item) || PySequence_Check(item)) {
         sequenceFound = true;
@@ -573,17 +575,25 @@ static bool THPTensor_(_checkAdvancedIndexing)(THPTensor *indexed, PyObject *arg
         Py_ssize_t start, end, length, step;
         if (THPUtils_parseSlice(item, dimSize, &start, &end, &step, &length)) {
           if (start != 0 || end != dimSize || step != 1 || length != dimSize) {
-            nonColonFound = true;
+            nonColonEllipsisFound = true;
             break;
           }
         }
         continue;
       }
-      nonColonFound = true;
+      if (Py_TYPE(item) == &PyEllipsis_Type) {
+        if (ellipsisFound) {
+          // Can't have duplicate ellipsi
+          return false;
+        }
+        ellipsisFound = true;
+        continue;
+      }
+      nonColonEllipsisFound = true;
       break;
     }
 
-    return sequenceFound && (!nonColonFound);
+    return sequenceFound && (!nonColonEllipsisFound);
   }
   return false;
 
@@ -639,6 +649,7 @@ static bool THPTensor_(_convertToTensorIndexers)(
   // 1. A LongTensor
   // 2. A sequence that can be converted into a LongTensor
   // 3. A empty slice object (i.e. ':')
+  // 4. An Ellipsis (i.e. '...')
   //
   // This function loops through all of the indexing elements. If we encounter
   // a LongTensor, we record the dimension at which it occurs. If we encounter
@@ -658,12 +669,29 @@ static bool THPTensor_(_convertToTensorIndexers)(
   std::vector<Py_ssize_t> indexingDims;
   std::vector<THPIndexTensor*>indexers;
 
+      // The indexing matches advanced indexing requirements. In the case that
+      // the user has an Ellipsis, and/or less dimensions than are in the
+      // Tensor being indexed, we "fill in" empty Slices to these dimensions
+      // so that the the resulting advanced indexing code still works
+
+
+
   // The top-level indexer should be a sequence, per the check above
   THPObjectPtr fast(PySequence_Fast(index, NULL));
   sequenceLength = PySequence_Fast_GET_SIZE(fast.get());
+  int ellipsisOffset = 0;
 
   for (Py_ssize_t i = 0; i < sequenceLength; ++i) {
     PyObject *item = PySequence_Fast_GET_ITEM(fast.get(), i);
+
+    // If this is an ellipsis, the all subsequent advanced indexing
+    // objects "positions" should be shifted, e.g. if we have a 5D Tensor
+    // x, and then x[..., [2, 3]], then the "position" of [2, 3] is 4
+    if (Py_TYPE(item) == &PyEllipsis_Type) {
+      ellipsisOffset = THTensor_(nDimension)(LIBRARY_STATE indexed) - sequenceLength;
+      continue;
+    }
+
     if (!PySlice_Check(item)) {
       // Returns NULL upon conversion failure
       THPIndexTensor *indexer = (THPIndexTensor *)PyObject_CallFunctionObjArgs(
@@ -680,7 +708,7 @@ static bool THPTensor_(_convertToTensorIndexers)(
         }
         return false;
       }
-      indexingDims.push_back(i);
+      indexingDims.push_back(i + ellipsisOffset);
       indexers.push_back(indexer);
     }
   }