Fix error message for cat-ing zero-dim tensors

zou3519 · zou3519 · commit 9ed3ebafdd29 · 2018-03-16T09:26:30.000-06:00
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -2305,7 +2305,9 @@ def test_cat_bad_input_sizes(self):
     def test_cat_scalars(self):
         x = torch.tensor(0)
         y = torch.tensor(1)
-        self.assertRaises(RuntimeError, lambda: torch.cat([x, y]))
+        with self.assertRaisesRegexp(RuntimeError,
+                                     'zero-dimensional.*cannot be concatenated'):
+            torch.cat([x, y])
 
     def test_stack(self):
         x = torch.rand(2, 3, 4)
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -155,7 +155,7 @@
   self: not_implemented("btrisolve")
 
 - name: cat(TensorList tensors, int64_t dim)
-  tensors: cat_tensors_backward(grad, to_arg_sizes(tensors, dim), dim)
+  tensors: cat_tensors_backward(grad, to_args_sizes(tensors), dim)
 
 - name: cauchy_(Tensor self, double median, double sigma, Generator generator)
   self: zeros_like(grad)
diff --git a/tools/autograd/load_derivatives.py b/tools/autograd/load_derivatives.py
@@ -294,10 +294,10 @@ def saved_variables(formula, args):
             'suffix': '_numel',
             'type': 'int64_t',
         }),
-        # replace to_arg_sizes(self, 2) with self_argsizes_2
-        (r'to_arg_sizes\({}, (\w+)\)', {
-            'suffix': lambda m: '_sizes_{}'.format(*m.groups()),
-            'type': 'IntList',
+        # replace to_args_sizes(self) with self_args_sizes
+        (r'to_args_sizes\({}\)', {
+            'suffix': '_args_sizes',
+            'type': 'std::vector<std::vector<int64_t>>',
         }),
         # replace TensorGeometry(self) with self_geometry
         (r'TensorGeometry\({}\)', {
diff --git a/tools/autograd/templates/Functions.cpp b/tools/autograd/templates/Functions.cpp
@@ -364,17 +364,23 @@ Tensor unsqueeze_to(const Tensor & self, int64_t dim, IntList sizes) {
   return self;
 }
 
-std::vector<Tensor> cat_tensors_backward(const Tensor & grad, const std::vector<int64_t> &sizes, int64_t dim) {
+std::vector<Tensor> cat_tensors_backward(const Tensor & grad, const std::vector<std::vector<int64_t>> &sizes, int64_t dim) {
+  if (sizes.size() > 0) {
+    // cat wraps dim to the first tensor's shape 
+    dim = at::maybe_wrap_dim(dim, sizes[0].size());
+  }
   std::vector<Tensor> grad_inputs(sizes.size());
   int64_t accumulate = 0;
   for (size_t i = 0; i < sizes.size(); ++i) {
-    auto size = sizes[i];
-    accumulate += size;
-    if (size == 0) {
+    auto& shape = sizes[i];
+    // If input was empty tensor, gradInput should be empty tensor.
+    if (shape[0] == 0) {
       grad_inputs[i] = at::zeros(grad.type(), {0});
-    } else {
-      grad_inputs[i] = grad.narrow(dim, accumulate - size, size);
+      continue;
     }
+    auto size = shape[dim];
+    accumulate += size;
+    grad_inputs[i] = grad.narrow(dim, accumulate - size, size);
   }
   return grad_inputs;
 }
diff --git a/tools/autograd/templates/VariableType.cpp b/tools/autograd/templates/VariableType.cpp
@@ -409,12 +409,12 @@ Tensor VariableType::contiguous(const Tensor & self) const {
   return self.clone();
 }
 
-static std::vector<int64_t> to_arg_sizes(TensorList tensors, int64_t dim) {
-  std::vector<int64_t> arg_sizes(tensors.size());
+static std::vector<std::vector<int64_t>> to_args_sizes(TensorList tensors) {
+  std::vector<std::vector<int64_t>> args_sizes(tensors.size());
   for (size_t i = 0; i < tensors.size(); ++i) {
-    arg_sizes[i] = tensors[i].size(dim);
+    args_sizes[i] = tensors[i].sizes();
   }
-  return arg_sizes;
+  return args_sizes;
 }
 
 ${type_derived_method_definitions}

Original file line number	Diff line number	Diff line change
`@@ -409,12 +409,12 @@ Tensor VariableType::contiguous(const Tensor & self) const {`
`409`	`409`	`return self.clone();`
`410`	`410`	`}`
`411`	`411`
`412`		`-static std::vector<int64_t> to_arg_sizes(TensorList tensors, int64_t dim) {`
`413`		`- std::vector<int64_t> arg_sizes(tensors.size());`
	`412`	`+static std::vector<std::vector<int64_t>> to_args_sizes(TensorList tensors) {`
	`413`	`+ std::vector<std::vector<int64_t>> args_sizes(tensors.size());`
`414`	`414`	`for (size_t i = 0; i < tensors.size(); ++i) {`
`415`		`- arg_sizes[i] = tensors[i].size(dim);`
	`415`	`+ args_sizes[i] = tensors[i].sizes();`
`416`	`416`	`}`
`417`		`- return arg_sizes;`
	`417`	`+ return args_sizes;`
`418`	`418`	`}`
`419`	`419`
`420`	`420`	`${type_derived_method_definitions}`