address comments

Roy Li · Roy Li · commit 2153c5eadad3 · 2018-06-18T14:38:03.000-07:00
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -600,21 +600,22 @@ Tensor & unsqueeze_(Tensor& self, int64_t dim) {
   return self.as_strided_(std::get<0>(g), std::get<1>(g));
 }
 
-Tensor flatten(const Tensor& self, int64_t start, int64_t end) {
-  start = maybe_wrap_dim(start, self.dim());
-  end = maybe_wrap_dim(end, self.dim());
-  AT_CHECK(start < end, "start dim must be before end dim");
+Tensor flatten(const Tensor& self, int64_t start_dim, int64_t end_dim) {
+  start_dim = maybe_wrap_dim(start_dim, self.dim());
+  end_dim = maybe_wrap_dim(end_dim, self.dim());
+  AT_CHECK(start_dim < end_dim, "start_dim must be before end_dim");
 
   std::vector<int64_t> shape;
-  for (int i = 0; i < start; i++) {
-    shape.push_back(self.sizes()[i]);
+  shape.reserve(self.dim() - end_dim + start_dim);
+  for (int64_t i = 0; i < start_dim; i++) {
+    shape.push_back(self.size(i));
   }
   shape.push_back(-1);
-  for (int i = end + 1; i < self.dim(); i++) {
-    shape.push_back(self.sizes()[i]);
+  for (int64_t i = end_dim + 1; i < self.dim(); i++) {
+    shape.push_back(self.size(i));
   }
 
-  return self.view(IntList(shape));
+  return self.reshape(shape);
 }
 
 Tensor view_as(const Tensor& self, const Tensor& other) {
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -556,7 +556,7 @@
   variants: function
   deprecated: true
 
-- func: flatten(Tensor self, int64_t start, int64_t end) -> Tensor
+- func: flatten(Tensor self, int64_t start_dim=0, int64_t end_dim=-1) -> Tensor
 
 - func: fill_(Tensor self, Scalar value) -> Tensor
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -5430,7 +5430,7 @@ def _fill_indices(self, idx, dim, dim_size, elems_per_row, m, n, o):
                     idx[tuple(ii)] = torch.randperm(dim_size)[0:elems_per_row]
 
     def test_flatten(self):
-        src = torch.randn(5,5,5,5)
+        src = torch.randn(5, 5, 5, 5)
         flat = src.flatten(0, -1)
         self.assertEqual(flat.shape, torch.Size([625]))
         self.assertEqual(src.view(-1), flat.view(-1))
@@ -5460,7 +5460,7 @@ def test_flatten(self):
             src.flatten(5, 10)
 
         # invalid start and end
-        with self.assertRaisesRegex(RuntimeError, 'start dim must be before end dim'):
+        with self.assertRaisesRegex(RuntimeError, 'start_dim must be before end_dim'):
             src.flatten(2, 0)
 
     @staticmethod
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -265,7 +265,7 @@
   self: zeros_like(grad)
   value: grad.sum()
 
-- name: flatten(Tensor self, int64_t start, int64_t end)
+- name: flatten(Tensor self, int64_t start_dim, int64_t end_dim)
   self: grad.view(self.sizes())
 
 - name: floor(Tensor self)
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -364,14 +364,14 @@ def reference_args(args):
 
     def get_trace_outputs(declaration):
         if declaration['return_type'] == 'std::vector<Tensor>':
-            return 'flatten_tensor({})'.format(declaration['returns'][0]['name'])
+            return 'flatten_tensor_args({})'.format(declaration['returns'][0]['name'])
         elif name.endswith('_out'):
             output_args = [arg['name'] for arg in arguments
                            if arg.get('output', False)]
             return '{' + ', '.join(output_args) + '}'
         trace_outs = [r['name'] for r in declaration['returns']]
         if any(ret['dynamic_type'] == 'TensorList' for ret in declaration['returns']):
-            return CodeTemplate("flatten_tensor( ${outs} )").substitute(outs=trace_outs)
+            return CodeTemplate("flatten_tensor_args( ${outs} )").substitute(outs=trace_outs)
         else:
             return CodeTemplate("{ ${outs} }").substitute(outs=trace_outs)
 
@@ -408,7 +408,7 @@ def emit_record_trace(env):
         local['tensor_args'] = [arg['name'] for arg in tensor_args]
         if any(arg['simple_type'] == 'TensorList' for arg in tensor_args):
             # Allocate a temporary vector with flatten and pass it in
-            local['trace_inputs'] = CodeTemplate("flatten_tensor( $tensor_args )").substitute(local)
+            local['trace_inputs'] = CodeTemplate("flatten_tensor_args( $tensor_args )").substitute(local)
         else:
             local['trace_inputs'] = CodeTemplate("{ ${tensor_args} }").substitute(local)
 
@@ -496,7 +496,7 @@ def emit_history():
         fn = 'rebase' if modifies_arguments and not is_view else 'set'
         output_names = [r['name'] for r in differentiable_outputs]
         # TODO: flatten allocates a std::vector, which could be expensive
-        outs = CodeTemplate("flatten_tensor( ${outs} )").substitute(outs=output_names)
+        outs = CodeTemplate("flatten_tensor_args( ${outs} )").substitute(outs=output_names)
         return SET_HISTORY.substitute(fn=fn, differentiable_outputs=outs)
 
     def emit_save_outputs():
diff --git a/tools/autograd/templates/VariableType.cpp b/tools/autograd/templates/VariableType.cpp
@@ -406,7 +406,7 @@ struct Flatten : IterArgs<Flatten> {
   }
 };
 
-template<typename... Args> inline variable_list flatten_tensor(Args&&... args) {
+template<typename... Args> inline variable_list flatten_tensor_args(Args&&... args) {
   variable_list out;
   out.reserve(count_tensors(std::forward<Args>(args)...));
   Flatten(out).apply(std::forward<Args>(args)...);
diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py
@@ -812,7 +812,7 @@ def add_docstr_all(method, docstr):
 
 add_docstr_all('flatten',
                r"""
-flatten(input, start, end) -> Tensor
+flatten(input, start_dim=0, end_dim=-1) -> Tensor
 
 see :func:`torch.flatten`
 """)
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -1593,14 +1593,26 @@ def parse_kwargs(desc):
 
 add_docstr(torch.flatten,
            r"""
-flatten(input, start, end) -> Tensor
+flatten(input, start_dim=0, end_dim=-1) -> Tensor
 
 Flattens a contiguous range of dims in a tensor.
 
 Args:
     input (Tensor): the input tensor
-    start (int): the first dim to flatten
-    end (int): the last dim to flatten
+    start_dim (int): the first dim to flatten
+    end_dim (int): the last dim to flatten
+
+Example::
+
+    >>> t = torch.tensor([[[1, 2],
+                           [3, 4]],
+                          [[5, 6],
+                           [7, 8]]])
+    >>> torch.flatten(t)
+    tensor([1, 2, 3, 4, 5, 6, 7, 8])
+    >>> torch.flatten(t, start_dim=1)
+    tensor([[1, 2, 3, 4],
+            [5, 6, 7, 8]])
 """)
 
 add_docstr(torch.gather,

Original file line number	Diff line number	Diff line change
`@@ -406,7 +406,7 @@ struct Flatten : IterArgs<Flatten> {`
`406`	`406`	`}`
`407`	`407`	`};`
`408`	`408`
`409`		`-template<typename... Args> inline variable_list flatten_tensor(Args&&... args) {`
	`409`	`+template<typename... Args> inline variable_list flatten_tensor_args(Args&&... args) {`
`410`	`410`	`variable_list out;`
`411`	`411`	`out.reserve(count_tensors(std::forward<Args>(args)...));`
`412`	`412`	`Flatten(out).apply(std::forward<Args>(args)...);`