Pytorch: smarter way to reduce dimension by reshape

Question

I want to reshape a Tensor by multiplying the shape of first two dimensions.

For example,

1st_tensor: torch.Size([12, 10]) to torch.Size([120])

2nd_tensor: torch.Size([12, 10, 5, 4]) to torch.Size([120, 5, 4])

I.e. The first two dimensions shall be merged into one, while the other dimensions shall remain the same.

Is there a smarter way than

1st_tensor.reshape(-1,)

2nd_tensor.reshape(-1,5,4),

that can adapt to the shape of different Tensors?

---

Test cases:

import torch
tests = [
    torch.rand(11, 11), 
    torch.rand(12, 15351, 6, 4),
    torch.rand(13, 65000, 8)
    ]

Answer 1

For tensor t, you can use:

t.reshape((-1,)+t.shape[2:])

This uses -1 for flattening the first two dimensions, and then uses t.shape[2:] to keep the other dimensions identical to the original tensor.

For your examples:

>>> tests = [
...     torch.rand(11, 11),
...     torch.rand(12, 15351, 6, 4),
...     torch.rand(13, 65000, 8)
...     ]
>>> tests[0].reshape((-1,)+tests[0].shape[2:]).shape
torch.Size([121])
>>> tests[1].reshape((-1,)+tests[1].shape[2:]).shape
torch.Size([184212, 6, 4])
>>> tests[2].reshape((-1,)+tests[2].shape[2:]).shape
torch.Size([845000, 8])

Answer 2

There is actually a better way I think:

import torch

tests = [
    torch.rand(11, 11), 
    torch.rand(12, 15351, 6, 4),
    torch.rand(13, 65000, 8)
    ]

for test in tests:
    print(test.flatten(0, 1).shape)

I'm not sure what version of torch this was added.