Repository: pytorch/pytorch

Issue:
[pytorch] [feature request] Flatten convenience method
Minor suggestion (trivial to implement in user code, but having it in the library would improve code brevity). The purpose is to flatten specific trailing dimensions by passing negative dimension index.

Can be useful for aggregating across multiple trailing dimensions, before mean/max etc get multiple dimensions support.

Exists in numpy/tensorflow/onnx, but semantics there doesn't allow flattening only specific dimensions.

```python
def flatten(x, dim):
    return x.view(x.size()[:dim] + (-1, ))

flatten(torch.rand(2,3,4,5,6), dim = -2).shape
# (2, 3, 4, 30)
```


Discussion:
Seems nice, could it also flatten multiple dimensions in a contiguous range? 
The snippet was just to flatten a few contiguous starting or trailing dimensions. If this is implemented into Core, maybe it’s better to rename `dim` to `multi_dims`? to align with recent multiple dimensions support in `sum`
Then maybe a naive support for `multi_dims` can be added to other ops like SoftMax by using `flatten` under the hood, at least if `multi_dims` contains contiguous dimensions (in actual memory layout).

Resolving patch:
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
--- 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({})'.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( ${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_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( ${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/torch/_tensor_docs.py b/torch/_tensor_docs.py
--- a/torch/_tensor_docs.py
+++ b/torch/_tensor_docs.py
@@ -810,6 +810,13 @@ def add_docstr_all(method, docstr):
 In-place version of :meth:`~Tensor.frac`
 """)
 
+add_docstr_all('flatten',
+               r"""
+flatten(input, start_dim=0, end_dim=-1) -> Tensor
+
+see :func:`torch.flatten`
+""")
+
 add_docstr_all('gather',
                r"""
 gather(dim, index) -> Tensor
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -1591,6 +1591,30 @@ def parse_kwargs(desc):
     array([-1,  2,  3])
 """)
 
+add_docstr(torch.flatten,
+           r"""
+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_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,
            r"""
 gather(input, dim, index, out=None) -> Tensor


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -5429,6 +5429,43 @@ def _fill_indices(self, idx, dim, dim_size, elems_per_row, m, n, o):
                     ii[dim] = slice(0, idx.size(dim) + 1)
                     idx[tuple(ii)] = torch.randperm(dim_size)[0:elems_per_row]
 
+    def test_flatten(self):
+        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))
+
+        flat = src.flatten(0, 2)
+        self.assertEqual(flat.shape, torch.Size([125, 5]))
+        self.assertEqual(src.view(-1), flat.view(-1))
+
+        flat = src.flatten(0, 1)
+        self.assertEqual(flat.shape, torch.Size([25, 5, 5]))
+        self.assertEqual(src.view(-1), flat.view(-1))
+
+        flat = src.flatten(1, 2)
+        self.assertEqual(flat.shape, torch.Size([5, 25, 5]))
+        self.assertEqual(src.view(-1), flat.view(-1))
+
+        flat = src.flatten(2, 3)
+        self.assertEqual(flat.shape, torch.Size([5, 5, 25]))
+        self.assertEqual(src.view(-1), flat.view(-1))
+
+        flat = src.flatten(-2, -1)
+        self.assertEqual(flat.shape, torch.Size([5, 5, 25]))
+        self.assertEqual(src.view(-1), flat.view(-1))
+
+        flat = src.flatten(2, 2)
+        self.assertEqual(flat, src)
+
+        # out of bounds index
+        with self.assertRaisesRegex(RuntimeError, 'dimension out of range'):
+            src.flatten(5, 10)
+
+        # invalid start and end
+        with self.assertRaisesRegex(RuntimeError, 'start_dim cannot come after end_dim'):
+            src.flatten(2, 0)
+
     @staticmethod
     def _test_gather(self, cast, test_bounds=True):
         m, n, o = random.randint(10, 20), random.randint(10, 20), random.randint(10, 20)