[FR] multidim squeeze and flatten

[ssnl]

squeeze

Previous FR for multidim squeeze: #1951

A PR (#4692) that attempted to implement #1951 was closed because

The behavior of this feature is very non-trivial, e.g., when negative and positive indices are mixed. I'm not sure it solves more problems that it adds anymore.

I would argue against that point, particularly for squeeze.

For squeeze, multidim is perfectly intuitive and defined, i.e., just squeezing the current dimensions the tuple refers. In fact, using a chain of squeeze (as suggested by the comment of #4692) is rather unintuitive.

If one has a tensor that has size 1 at dims 2 and 3, which they want to squeeze, it is far more intuitive to write

t.squeeze((2, 3))

than to write

t.squeeze(2).squeeze(2)

Negative indices does not add to complexity or confusion either.

flatten

Similarly, to merge the dims 1&2 and dims 3&4, it is far more intuitive to write

t.flatten([(1, 2), (3, 4)])

than

t.flatten(1,2).flatten(2,3)

[zou3519]

I agree that it makes sense for squeeze and flatten to have these APIs. I would be happy to review PRs that enable this.

[vadimkantorov]

squeeze_multiple already exists somewhere in PyTorch

unsqueeze_multiple also exists

another instance of real-world unsqueeze_multiple:

pytorch/torch/csrc/jit/symbolic_script.cpp

Line 45 in c4f10e0

grad = AD_unsqueeze_multiple(grad, dim, n_dims)

[vadimkantorov]

I recently had a practical usecase for unsqueeze_multiple, had to implement it like this:
unsqueeze_multiple = lambda x, dim, k: functools.reduce(lambda acc, _: acc.unsqueeze(dim), range(k), x) (https://gist.github.com/vadimkantorov/ba3d3d099e403f2703facd2765625f27)

[ssnl]

Oh, so for unsqueeze_multiple, it is just inserting multiple dims at the same location, right? @vadimkantorov

[vadimkantorov]

Yup! Maybe this version better called unsqueeze_repeated. It was useful for forcing a tensor to have a given number of dimensions e.g. 4d

[vadimkantorov]

A recent related example: #94750

The code in that example is impl of mutlidim squeeze: x.squeeze(dim = (d1, d2, ...))

[vadimkantorov]

It appears that multidim squeeze is merged now :)