Make cat/cat_out native function that rejects scalar inputs. (#4992)

gchanan · web-flow · commit eb5daa94780c · 2018-02-01T16:20:01.000-05:00
* Make cat/cat_out native function that rejects scalar inputs.

* Print position of scalar in error message.
diff --git a/aten/src/ATen/Declarations.cwrap b/aten/src/ATen/Declarations.cwrap
@@ -4162,7 +4162,7 @@
 ]]
 
 [[
-  name: cat
+  name: _cat
   cname: catArray
   variants: [function]
   return: self
diff --git a/aten/src/ATen/native/Convolution.cpp b/aten/src/ATen/native/Convolution.cpp
@@ -382,7 +382,7 @@ at::Tensor _convolution(
         outputs[g] = at::_convolution_nogroup(
             input_g, weight_g, bias_g, params.stride, params.padding, params.dilation, params.transposed, params.output_padding);
       }
-      output = cat(outputs, 1);
+      output = at::cat(outputs, 1);
     }
   }
 
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -8,6 +8,25 @@
 namespace at {
 namespace native {
 
+static void check_cat_no_zero_dim(TensorList tensors) {
+  for(size_t i = 0; i < tensors.size(); ++i) {
+    auto& t = tensors[i];
+    if (t.dim() == 0) {
+      runtime_error("zero-dimensional tensor (at position %zu) cannot be concatenated", i);
+    }
+  }
+}
+
+Tensor & cat_out(Tensor & result, TensorList tensors, int64_t dim) {
+  check_cat_no_zero_dim(tensors);
+  return at::_cat_out(result, tensors, dim);
+}
+
+Tensor cat(TensorList tensors, int64_t dim) {
+  check_cat_no_zero_dim(tensors);
+  return at::_cat(tensors, dim);
+}
+
 std::vector<Tensor> chunk(const Tensor& self, int64_t chunks, int64_t dim) {
   if (self.dim() == 0) {
     throw std::runtime_error("chunk expects at least a 1-dimensional tensor");
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -47,6 +47,12 @@
 
 - func: bernoulli_(Tensor self, double p=0.5, Generator* generator=nullptr) -> Tensor
 
+- func: cat(TensorList tensors, int64_t dim=0) -> Tensor
+  variants: function
+
+- func: cat_out(Tensor result, TensorList tensors, int64_t dim=0) -> Tensor
+  variants: function
+
 - func: chunk(Tensor self, int64_t chunks, int64_t dim=0) -> TensorList
 
 - func: cudnn_is_acceptable(Tensor self) -> bool
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -1966,6 +1966,13 @@ def test_cat_bad_input_sizes(self):
         z = torch.randn(2, 2, 1)
         self.assertRaises(RuntimeError, lambda: torch.cat([x, y, z], dim=1))
 
+    @unittest.skipIf(not torch._C._with_scalars(), "scalars not enabled")
+    def test_cat_scalars(self):
+        from torch.autograd import variable
+        x = variable(0)
+        y = variable(1)
+        self.assertRaises(RuntimeError, lambda: torch.cat([x, y]))
+
     def test_stack(self):
         x = torch.rand(2, 3, 4)
         y = torch.rand(2, 3, 4)

Original file line number	Diff line number	Diff line change
`@@ -382,7 +382,7 @@ at::Tensor _convolution(`
`382`	`382`	`outputs[g] = at::_convolution_nogroup(`
`383`	`383`	`input_g, weight_g, bias_g, params.stride, params.padding, params.dilation, params.transposed, params.output_padding);`
`384`	`384`	`}`
`385`		`- output = cat(outputs, 1);`
	`385`	`+ output = at::cat(outputs, 1);`
`386`	`386`	`}`
`387`	`387`	`}`
`388`	`388`