Unify error checking for tesnor.index_copy_ (pytorch#5642)

Author: zou3519

aten/src/ATen/Declarations.cwrap

@@ -258,8 +258,7 @@
     - THIndexTensor* index
 ]]
 [[
-  name: indexCopy_
-  python_name: index_copy_
+  name: _indexCopy_
   cname: indexCopy
   return: argument 0
   arguments:

aten/src/ATen/native/Indexing.cpp

@@ -250,4 +250,49 @@ Tensor & index_put_(Tensor & self, TensorList indices, const Tensor & value) {
   return src.put_(linearIndex, expandedValue);
 }
 
+Tensor & index_copy_(Tensor & self, int64_t dim, const Tensor & index, const Tensor & source) {
+  dim = maybe_wrap_dim(dim, self.dim());
+
+  if (index.dim() >= 2) {
+    runtime_error(
+        "index_copy_(): Index should have dimension 1 or 0 (got %d)",
+        (int)index.dim());
+  }
+  int64_t numIndices = index.numel();
+  if (source.dim() == 0 && numIndices != 1) {
+    runtime_error(
+        "index_copy_(): When source is scalar, index should have one element (got %d)",
+        (int)numIndices);
+  }
+  if (source.dim() > 0 && numIndices != source.size(dim)) {
+    runtime_error(
+        "index_copy_(): Number of indices (%d) should be equal to source.size(dim) (%d)",
+        (int)numIndices, (int)source.size(dim));
+  }
+  if (index.type().scalarType() != ScalarType::Long) {
+    runtime_error("index_copy_(): Expected LongTensor for index");
+  }
+
+  // Check that source and destination slices have the same size
+  auto selfSlicedSizes = std::vector<int64_t>(self.sizes());
+  if (selfSlicedSizes.size() > 0) {
+    selfSlicedSizes.erase(selfSlicedSizes.begin() + dim);
+  }
+  auto sourceSlicedSizes = std::vector<int64_t>(source.sizes());
+  if (sourceSlicedSizes.size() > 0) {
+    sourceSlicedSizes.erase(sourceSlicedSizes.begin());
+  }
+  if (selfSlicedSizes.size() != sourceSlicedSizes.size() ||
+      !std::equal(selfSlicedSizes.begin(), selfSlicedSizes.end(),
+                  sourceSlicedSizes.begin())) {
+    std::stringstream ss;
+    ss << "index_copy_(): Source/destination tensor must have same slice shapes. ";
+    ss << "Destination slice shape: " << selfSlicedSizes << " at dimension " << dim;
+    ss << " and source slice shape: " << sourceSlicedSizes << " at dimension 0.";
+    throw std::runtime_error(ss.str());
+  }
+
+  return self._indexCopy_(dim, index, source);
+}
+
 }} // at::native

aten/src/ATen/native/native_functions.yaml

@@ -294,6 +294,9 @@
 - func: index(Tensor self, TensorList indices) -> Tensor
   # NB: This function is special-cased in tools/autograd/gen_variable_type.py
 
+- func: index_copy_(Tensor self, int64_t dim, IndexTensor index, Tensor source) -> Tensor
+  variants: method
+
 - func: index_put_(Tensor self, TensorList indices, Tensor values) -> Tensor
 
 - func: is_cuda(Tensor self) -> bool

aten/src/TH/generic/THTensorMath.c

@@ -382,10 +382,9 @@ void THTensor_(indexCopy)(THTensor *tensor, int dim, THLongTensor *index, THTens
   THTensor *tSlice, *sSlice;
   int64_t *index_data;
 
+  // Error checking for this function has moved to ATen!!
+
   numel = THLongTensor_nElement(index);
-  THArgCheck(index->nDimension == 1, 3, "Index is supposed to be a vector");
-  THArgCheck(dim < src->nDimension, 4, "Indexing dim %d is out of bounds of tensor", dim + TH_INDEX_BASE);
-  THArgCheck(numel == src->size[dim],4,"Number of indices should be equal to source:size(dim)");
 
   index = THLongTensor_newContiguous(index);
   index_data = THLongTensor_data(index);

tools/autograd/gen_python_functions.py

@@ -17,7 +17,7 @@
     'alias', 'contiguous', 'clamp.*', 'is_cuda', 'is_sparse', 'size', 'stride',
     '.*_backward', '.*_backward_out', '.*_forward', '.*_forward_out',
     'sparse_raw_resize_', '_unsafe_view', 'tensor', 'sparse_coo_tensor',
-    '_arange.*', '_range.*', '_linspace.*', '_logspace.*'
+    '_arange.*', '_range.*', '_linspace.*', '_logspace.*', '_indexCopy_',
 ]
 
 PY_VARIABLE_METHODS_CPP = CodeTemplate.from_file(template_path + '/python_variable_methods.cpp')