[Feature suggestion] splatting list of anonymous dimensions of any length

[lucidrains]

Hi Alex again! 👋 Thank you again for this wonderful work 🙏 I still can't get over how beautiful the abstraction is, and how much of a multiplier it is in my work.

I was wondering what your thoughts are about this potential extension to einops? A common pattern I run into is the need to flatten dimensions (for some operation), and then to reconstitute it later. To do this, I usually need to save the dimensions to a list or tuple, and then later run a reshape on the output. To do this with einops rearrange currently is difficult, as it is unable to take in a list of dimensions (it needs each dimension to be explicitly named), and there is no way to name the ellipsis-es.

I hacked up a prototype here lucidrains/memorizing-transformers-pytorch@b190c5e to give you an idea of the use-case. Would be curious to hear your thoughts, whether you can think of some even more general way to handle this with the internal ellipsis code you have already, or whether you think this doesn't belong in einops proper

Thank you!

[arogozhnikov]

Thank you, always glad to hear that einops is helpful.

I hacked up a prototype here lucidrains/memorizing-transformers-pytorch@b190c5e to give you an idea of the use-case.

What you suggest with named ellipsis is exactly how I was thinking about it.

This pathway has some obstacles:

1. it's unclear how scripting / tracing would work for backends if passed arguments can be integers, symbols, tuples of integers, tuples of symbols or symbol-tuple, this will require much testing. Scripters are usually not that flexible.

2. Another consideration is layers - they have no way to pass / receive shapes on the go.

3. Should this feature appear, parse shaping should be also extended - in this case output types would be integer or tuple.

That said, I'd like to have named ellipses in the notation.

A common pattern I run into is the need to flatten dimensions (for some operation), and then to reconstitute it later.

I think generally common use-case is "to (partially) reconstruct structure after the operation" and flattening is one of cases.
When concatenating tensors, users frequently want to split back afterwards output tensor in the same way input was constructed (that's one of things that I want to have properly designed and aligned before moving further).

There are some simple designs with objects-keepers, but they can't provide caching and scripting at the same time.

[lucidrains]

@arogozhnikov whatever you think is best Alex! 🙏 just wanted to incept the idea :)

[arogozhnikov]

Hey @lucidrains, I believe einops.pack and einops.unpack can do what you ask for

https://github.com/arogozhnikov/einops/blob/master/docs/4-pack-and-unpack.ipynb

[lucidrains]

@arogozhnikov omg! this is brilliant! 🙏

[lucidrains]

can't wait for 0.6!