Fix indexing when passing only an Ellipsis

Author: fmassa

test/test_torch.py

@@ -1880,11 +1880,13 @@ def test_index(self):
  self.assertEqual(reference[2, ..., 2, 2], 27, 0)
  self.assertEqual(reference[2, 2, ..., 2], 27, 0)
  self.assertEqual(reference[2, 2, 2, ...], 27, 0)
+ self.assertEqual(reference[...], reference, 0)
 
  reference_5d = self._consecutive((3, 3, 3, 3, 3))
  self.assertEqual(reference_5d[..., 1, 0], reference_5d[:, :, :, 1, 0], 0)
  self.assertEqual(reference_5d[2, ..., 1, 0], reference_5d[2, :, :, 1, 0], 0)
  self.assertEqual(reference_5d[2, 1, 0, ..., 1], reference_5d[2, 1, 0, :, 1], 0)
+ self.assertEqual(reference_5d[...], reference_5d, 0)
 
  # LongTensor indexing
  reference = self._consecutive((5, 5, 5))

torch/csrc/generic/Tensor.cpp

@@ -540,7 +540,8 @@ static bool THPTensor_(_index)(THPTensor *self, PyObject *index,
  }
  }
  if (valid) return true;
- } else {
+ } else if (index == Py_Ellipsis) return true;
+ else {
  if (THPTensor_(_indexOnce)<allow_index>(index, indexed_dim, tresult, sresult, storage_offset))
  return true;
  }
@@ -551,9 +552,9 @@ static bool THPTensor_(_index)(THPTensor *self, PyObject *index,
  ", numpy scalars"
 #endif
 #ifndef THC_GENERIC_FILE
- "torch.LongTensor and torch.ByteTensor.",
+ ", torch.LongTensor and torch.ByteTensor.",
 #else
- "torch.cuda.LongTensor and torch.cuda.ByteTensor.",
+ ", torch.cuda.LongTensor and torch.cuda.ByteTensor.",
 #endif
  THPUtils_typename(index));
  return false;